diff --git a/nemo_automodel/recipes/diffusion/train.py b/nemo_automodel/recipes/diffusion/train.py index 146fd9ba51..20b0ae6d82 100644 --- a/nemo_automodel/recipes/diffusion/train.py +++ b/nemo_automodel/recipes/diffusion/train.py @@ -232,6 +232,12 @@ def _build_diffusion_parallel_manager_args( if compute_dtype is None: compute_dtype = dtype + # The recipe passes ConfigNode sections, which support .to_dict() but not dict(). + if hasattr(fsdp_cfg, "to_dict"): + fsdp_cfg = fsdp_cfg.to_dict() + if hasattr(ddp_cfg, "to_dict"): + ddp_cfg = ddp_cfg.to_dict() + if fsdp_cfg is not None and ddp_cfg is not None: raise ValueError( "Cannot specify both 'fsdp' and 'ddp' configurations. " diff --git a/tests/ci_tests/scripts/diffusion_finetune_launcher.sh b/tests/ci_tests/scripts/diffusion_finetune_launcher.sh index c25f8f1561..8806ca326a 100755 --- a/tests/ci_tests/scripts/diffusion_finetune_launcher.sh +++ b/tests/ci_tests/scripts/diffusion_finetune_launcher.sh @@ -144,13 +144,21 @@ fi echo "============================================" echo "[finetune] Running finetuning..." echo "============================================" +# The recipe rejects configs that contain both 'fsdp' and 'ddp' sections, and +# a --fsdp.* CLI override injects an 'fsdp' section. Only pass it for FSDP +# recipes; DDP replicates across all ranks and needs no dp_size. +DIST_OVERRIDE="--fsdp.dp_size ${NPROC_PER_NODE}" +if grep -qE '^ddp:' "/opt/Automodel/${CONFIG_PATH}"; then + DIST_OVERRIDE="" +fi + CONFIG="--config /opt/Automodel/${CONFIG_PATH} \ --data.dataloader.cache_dir $DATA_DIR/cache \ --checkpoint.checkpoint_dir $CKPT_DIR \ --step_scheduler.max_steps ${MAX_STEPS:-100} \ --step_scheduler.ckpt_every_steps 100 \ --step_scheduler.save_checkpoint_every_epoch false \ - --fsdp.dp_size ${NPROC_PER_NODE} \ + ${DIST_OVERRIDE} \ --wandb.mode disabled" CMD="uv run --extra diffusion torchrun --nproc-per-node=${NPROC_PER_NODE} \ diff --git a/tests/unit_tests/recipes/test_diffusion_train_metrics.py b/tests/unit_tests/recipes/test_diffusion_train_metrics.py index c0f8b3357c..5c424ed0ea 100644 --- a/tests/unit_tests/recipes/test_diffusion_train_metrics.py +++ b/tests/unit_tests/recipes/test_diffusion_train_metrics.py @@ -464,6 +464,36 @@ def test_build_diffusion_parallel_manager_args_parses_ddp_config(): } +def test_build_diffusion_parallel_manager_args_accepts_confignode_fsdp_config(): + manager_args = _build_diffusion_parallel_manager_args( + fsdp_cfg=ConfigNode({"dp_size": 8, "cpu_offload": False}), + ddp_cfg=None, + world_size=8, + dtype=torch.bfloat16, + lora_enabled=False, + ) + + assert manager_args["_manager_type"] == "fsdp2" + assert manager_args["dp_size"] == 8 + + +def test_build_diffusion_parallel_manager_args_accepts_confignode_ddp_config(): + manager_args = _build_diffusion_parallel_manager_args( + fsdp_cfg=None, + ddp_cfg=ConfigNode({"backend": "nccl", "activation_checkpointing": False}), + world_size=4, + dtype=torch.bfloat16, + lora_enabled=False, + ) + + assert manager_args == { + "_manager_type": "ddp", + "world_size": 4, + "activation_checkpointing": False, + "find_unused_parameters": False, + } + + def test_build_model_and_optimizer_forwards_perf_options_and_optimizer_kwargs(monkeypatch): pipe = SimpleNamespace(transformer=_TinyTransformer()) manager = SimpleNamespace(device_mesh="mesh")