update use_custom_fsdp to use_megatron_fsdp

Author: shjwudp

nemo/collections/diffusion/recipes/flux_12b.py

@@ -113,7 +113,7 @@ def trainer(
         gradient_accumulation_fusion=True,
         ddp=run.Config(
             DistributedDataParallelConfig,
-            # use_custom_fsdp=True,
+            # use_megatron_fsdp=True,
             # data_parallel_sharding_strategy='optim_grads_params',
             check_for_nan_in_grad=True,
             grad_reduce_in_fp32=True,

scripts/dit/dit_train.py

@@ -210,7 +210,7 @@ def train_mock() -> run.Partial:
     recipe.data.model_config = recipe.model.config
     recipe.log.log_dir = 'nemo_experiments/train_mock'
 
-    recipe.trainer.strategy.ddp.use_custom_fsdp = True
+    recipe.trainer.strategy.ddp.use_megatron_fsdp = True
     recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = 'optim_grads_params'
     recipe.trainer.strategy.ddp.overlap_param_gather = True
     recipe.trainer.strategy.ddp.overlap_grad_reduce = True
@@ -236,7 +236,7 @@ def mock_ditllama5b_8k() -> run.Partial:
     recipe.data.model_config = recipe.model.config
     recipe.log.log_dir = 'nemo_experiments/mock_ditllama5b_8k'
     recipe.model.config.attn_mask_type = AttnMaskType.no_mask
-    recipe.trainer.strategy.ddp.use_custom_fsdp = True
+    recipe.trainer.strategy.ddp.use_megatron_fsdp = True
     recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = 'optim_grads_params'
     recipe.trainer.strategy.ddp.overlap_param_gather = True
     recipe.trainer.strategy.ddp.overlap_grad_reduce = True
@@ -360,7 +360,7 @@ def pretrain_ditllama30b() -> run.Partial:
     recipe.data.task_encoder.seq_length = 256
     recipe.data.virtual_epoch_length = 0
     recipe.log.log_dir = 'nemo_experiments/ditllama30b_stage1_mock'
-    recipe.trainer.strategy.ddp.use_custom_fsdp = True
+    recipe.trainer.strategy.ddp.use_megatron_fsdp = True
     recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = 'optim_grads_params'
     recipe.trainer.strategy.ddp.overlap_param_gather = True
     recipe.trainer.strategy.ddp.overlap_grad_reduce = True
@@ -386,7 +386,7 @@ def pretrain_ditllama30b_stage2_mock() -> run.Partial:
     recipe.trainer.val_check_interval = 1.0
     recipe.data.model_config = recipe.model.config
     recipe.log.log_dir = 'nemo_experiments/ditllama30b_stage2_mock'
-    recipe.trainer.strategy.ddp.use_custom_fsdp = True
+    recipe.trainer.strategy.ddp.use_megatron_fsdp = True
     recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = 'optim_grads_params'
     recipe.trainer.strategy.ddp.overlap_param_gather = True
     recipe.trainer.strategy.ddp.overlap_grad_reduce = True
@@ -412,7 +412,7 @@ def pretrain_ditllama30b_stage3_mock() -> run.Partial:
     recipe.trainer.val_check_interval = 1.0
     recipe.data.model_config = recipe.model.config
     recipe.log.log_dir = 'nemo_experiments/ditllama30b_stage3_mock'
-    recipe.trainer.strategy.ddp.use_custom_fsdp = True
+    recipe.trainer.strategy.ddp.use_megatron_fsdp = True
     recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = 'optim_grads_params'
     recipe.trainer.strategy.ddp.overlap_param_gather = True
     recipe.trainer.strategy.ddp.overlap_grad_reduce = True
@@ -512,7 +512,7 @@ def pretrain_ecditllama1b() -> run.Partial:
     recipe.log.log_dir = 'nemo_experiments/ecditllama1b'
     recipe.trainer.val_check_interval = 3000
 
-    recipe.trainer.strategy.ddp.use_custom_fsdp = True
+    recipe.trainer.strategy.ddp.use_megatron_fsdp = True
     recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = 'optim_grads_params'
     recipe.trainer.strategy.ddp.overlap_param_gather = True
     recipe.trainer.strategy.ddp.overlap_grad_reduce = True

scripts/flux/flux_controlnet_training.py

@@ -92,7 +92,7 @@ def flux_controlnet_training() -> run.Partial:
                 pipeline_dtype=torch.bfloat16,
                 ddp=run.Config(
                     DistributedDataParallelConfig,
-                    use_custom_fsdp=True,
+                    use_megatron_fsdp=True,
                     data_parallel_sharding_strategy='optim_grads_params',
                     check_for_nan_in_grad=True,
                     grad_reduce_in_fp32=True,
@@ -292,7 +292,7 @@ def unit_test(custom_fsdp=True) -> run.Partial:
 def configure_custom_fsdp(recipe) -> run.Partial:
     recipe.trainer.strategy.ddp = run.Config(
         DistributedDataParallelConfig,
-        use_custom_fsdp=True,
+        use_megatron_fsdp=True,
         data_parallel_sharding_strategy='optim_grads_params',  # Custom FSDP
         check_for_nan_in_grad=True,
         grad_reduce_in_fp32=True,

scripts/flux/flux_training.py

@@ -95,7 +95,7 @@ def flux_training() -> run.Partial:
                 gradient_accumulation_fusion=True,
                 ddp=run.Config(
                     DistributedDataParallelConfig,
-                    use_custom_fsdp=True,
+                    use_megatron_fsdp=True,
                     data_parallel_sharding_strategy='optim_grads_params',
                     check_for_nan_in_grad=True,
                     grad_reduce_in_fp32=True,
@@ -229,7 +229,7 @@ def fp8_test(custom_fsdp=True) -> run.Partial:
 def configure_custom_fsdp(recipe) -> run.Partial:
     recipe.trainer.strategy.ddp = run.Config(
         DistributedDataParallelConfig,
-        use_custom_fsdp=True,
+        use_megatron_fsdp=True,
         data_parallel_sharding_strategy='optim_grads_params',  # Custom FSDP
         check_for_nan_in_grad=True,
         grad_reduce_in_fp32=True,

scripts/performance/llm/pretrain_llama3_8b.py

@@ -84,6 +84,49 @@ def override_recipe_configs(
     recipe = set_exp_logging_configs(
         recipe, "pre_train", "llm", "llama3", args.tensorboard, args.wandb, args.wandb_prj_name, args.wandb_job_name
     )
+    # for saving checkpoints
+    ckpt_path = "/lustre/fsw/coreai_devtech_all/jianbinc/playground/nemo_nvfsdp_update/NeMo/checkpoints"
+    recipe.log.log_dir = ckpt_path
+    import nemo.lightning as nl
+    import nemo_run as run
+
+    recipe.log.ckpt = run.Config(
+        nl.ModelCheckpoint,
+        train_time_interval=None,
+        save_last=True,
+        every_n_train_steps=100,
+        save_top_k=1,
+        save_on_train_epoch_end=True,
+        save_optim_on_train_end=True,
+        always_save_context=False,
+        filename="{model_name}--{val_loss:.2f}-{step}-{consumed_samples}",
+    )
+    
+    # nl.ModelCheckpoint(
+    #     train_time_interval=None,
+    # )
+    # # recipe.log.ckpt.train_time_interval = None
+    # recipe.log.ckpt.save_last = True
+    # recipe.log.ckpt.every_n_train_steps = 100
+    # recipe.log.ckpt.save_top_k = 1
+    # recipe.log.ckpt.save_on_train_epoch_end = True
+    # recipe.log.ckpt.save_optim_on_train_end = True
+    # recipe.log.ckpt.always_save_context = False
+
+    # for loading checkpoints
+    recipe.resume.resume_if_exists = True
+    recipe.resume.resume_ignore_no_checkpoint = True
+    # recipe.resume.restore_config = RestoreConfig(
+    #     path=ckpt_path,
+    #     load_model_state=True,
+    #     load_optim_state=True,
+    # )
+
+    recipe.trainer.strategy.save_ckpt_format = "fsdp_dtensor"
+    recipe.trainer.strategy.ddp.average_in_collective = False
+    # recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = "optim"
+
+    recipe.optim.config.use_precision_aware_optimizer = False
 
     # data module configs
     if args.use_hf_tokenizer: