Skip to content

Commit 70a9830

Browse files
committed
fix(recipes): prevent Qwen3.5-VL FSDP SFT OOM
Signed-off-by: Chen Cui <chcui@nvidia.com>
1 parent 8a84d71 commit 70a9830

2 files changed

Lines changed: 14 additions & 4 deletions

File tree

‎src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1190,7 +1190,9 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer:
11901190
cfg.model.attention_backend = "auto"
11911191
cfg.model.gradient_accumulation_fusion = True
11921192
cfg.model.cross_entropy_loss_fusion = True
1193-
cfg.model.cross_entropy_fusion_impl = "native"
1193+
# Native fused cross entropy materializes a full FP32 vocabulary-sized
1194+
# softmax buffer. Use the TE kernel to keep the 248K-token loss bounded.
1195+
cfg.model.cross_entropy_fusion_impl = "te"
11941196

11951197
# MoE settings
11961198
cfg.model.moe_token_dispatcher_type = "alltoall"
@@ -1204,10 +1206,12 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer:
12041206
cfg.model.moe_router_padding_for_fp8 = False
12051207

12061208
# Memory saving
1207-
cfg.model.recompute_granularity = None
1209+
# Full SFT retains FP32 optimizer state alongside the FSDP communication
1210+
# pools, so recompute each transformer layer to bound live activations.
1211+
cfg.model.recompute_granularity = "full"
12081212
cfg.model.recompute_modules = None
1209-
cfg.model.recompute_method = None
1210-
cfg.model.recompute_num_layers = None
1213+
cfg.model.recompute_method = "uniform"
1214+
cfg.model.recompute_num_layers = 1
12111215
cfg.model.fine_grained_activation_offloading = False
12121216
cfg.model.offload_modules = None
12131217

‎tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -561,6 +561,12 @@ def test_qwen35_vl_35b_a3b_fsdp_sft_defaults(monkeypatch: pytest.MonkeyPatch):
561561
assert cfg.model.sequence_parallel is False
562562
assert cfg.model.moe_token_dispatcher_type == "alltoall"
563563
assert cfg.model.moe_router_fusion is True
564+
assert cfg.model.cross_entropy_loss_fusion is True
565+
assert cfg.model.cross_entropy_fusion_impl == "te"
566+
assert cfg.model.recompute_granularity == "full"
567+
assert cfg.model.recompute_modules is None
568+
assert cfg.model.recompute_method == "uniform"
569+
assert cfg.model.recompute_num_layers == 1
564570
assert cfg.ddp.use_megatron_fsdp is True
565571
assert cfg.ddp.fsdp_double_buffer is True
566572
assert cfg.ddp.megatron_fsdp_max_pool_double_buffer is True

0 commit comments

Comments
 (0)