refactor: simplify training and inference interfaces

- avoid constructing model_fn more than once when reading config - keep inference package exports focused on public entry points - rename extra strategy arguments to strategy_kwargs
2026-08-19 20:55:13 +08:00
parent 398e8a3ea3
commit c79d34eee1
9 changed files with 17 additions and 83 deletions
@@ -836,7 +836,7 @@ def train(
        gradient_checkpointing_modules=grad_ckpt_modules,
        compile_mode=compile_mode,
        executor_kwargs=executor_kwargs,
-        extra_kwargs=strategy_kwargs,
+        strategy_kwargs=strategy_kwargs,
        neftune_alpha=neftune_alpha,
        collate_fn=collate_fn,
        rollout_interval=rollout_interval,