cp: fix(diffusion): resolve flux nightly CI failures (2529) into r0.5.0

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. "

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} \

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")