support num_workers,save_steps,find_unused_parameters

2026-03-22 16:50:47 +00:00 · 2025-08-06 10:52:59 +08:00
parent 8d2f6ad32e
commit 6bae70eee0
13 changed files with 71 additions and 16 deletions
--- a/examples/flux/model_training/train.py
+++ b/examples/flux/model_training/train.py
@@ -121,4 +121,7 @@ if __name__ == "__main__":
        dataset, model, model_logger, optimizer, scheduler,
        num_epochs=args.num_epochs,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        save_steps=args.save_steps,
+        find_unused_parameters=args.find_unused_parameters,
+        num_workers=args.dataset_num_workers,
    )