Fix OOM for gemma4 31B on tot for gemma4_31b.yaml

[athitten]

Describe the bug

After #1904 got merged, main is running into OOM while running gemma4_31b.yaml which uses just fsdp + ac (does not use any parallelism). Creating this issue to investigate this further.

Steps/Code to reproduce bug

Please list minimal steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

Expected behavior

A clear and concise description of what you expected to happen.

Additional context

Add any other context about the problem here.

[khazic]

For running Gemma4 31B, you can use the TP+PP yaml configs added in #1904 as a workaround:

• Single node (8 GPUs, TP4×PP2): examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml
• 4 nodes (32 GPUs, TP4×PP4×DP2): examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml

These should fit in memory where the plain FSDP config OOMs.

[khazic]

Looking into this — the OOM may have been introduced by my recent merged PR. Will investigate the root cause.

[khazic]

I have a hypothesis on the root cause. In PR #1904, stateful_wrappers.py changed the tied lm_head logic:

• Before: is_tied_lm_head was derived from the model config, so lm_head.weight was always popped from the state dict when tied embeddings were configured.
• After: has_local_tied_lm_head checks actual storage pointer sharing. FSDP2 sharding breaks the storage tie between lm_head.weight and embed_tokens.weight, so has_local_tied_lm_head returns False — meaning lm_head.weight is no longer dropped from the state dict.

For Gemma4 31B, the embedding matrix is ~3.5 GB (vocab_size 262144 × hidden_dim 7168 × bf16). If both embed_tokens and lm_head are fully gathered during checkpoint/validation (at step 500), that extra 3.5 GB per rank could push an already memory-tight FSDP run into OOM.

To confirm: check whether the OOM happens at step 1 (training forward/backward) or at step 500 (checkpoint + validation). If it's step 500, this is likely the cause.

I'll have a confirmed root cause and fix within 3 hours.

[khazic]

Root cause confirmed and fix submitted in #1941.

The regression was introduced by #1904: _reduce_attrs returns each nn.ModuleList as a single item, so layers = [ModuleList] instead of individual decoder layers. AC code then tries to find self_attn/mlp on the ModuleList object, finds nothing, and silently skips all checkpointing — causing OOM during training forward pass.

Fix: flatten nn.ModuleList results in _extract_model_layers so individual layers are returned, matching the heuristic path behavior.

[khazic]

Full root cause analysis, documented in #1941:

Why did #1904 (TP+PP) break AC without touching AC code?

_extract_model_layers serves dual purpose — both TP sharding and AC consume the same returned layers list:

layers = _extract_model_layers(model)

if tp_mesh.size() > 1:
    for layer in layers:
        parallelize_module(layer, tp_mesh, tp_plan)   # TP uses it

if activation_checkpointing:
    for layer in layers:
        apply_activation_checkpointing(layer, ...)    # AC uses it

To generate a TP sharding plan for Gemma4, #1904 registered it in MODEL_CLS_TO_LAYERS. This moved Gemma4 from the heuristic path (which worked correctly) to the registered path (which has a bug).

Before #1904 — Gemma4 used the heuristic path, and plain FSDP2 + AC on 31B worked fine:

elif hasattr(model, "model") and hasattr(model.model, "layers"):
    layers.extend(model.model.layers)   # iterates ModuleList directly → individual layers ✓

After #1904 — Gemma4 hits the registered path:

layers.extend(_reduce_attrs(model, ["model.language_model.layers"]))
# _reduce_attrs returns [<nn.ModuleList>] — one item, the whole ModuleList
# layers = [ModuleList] instead of [layer0, layer1, ..., layer61]
# AC code looks for self_attn/mlp on ModuleList → not found → silently skips all 62 layers → OOM

Is the bug in the registered path pre-existing?

The same layers.extend(_reduce_attrs(...)) pattern was already present for all other registered VLMs (Qwen2VL, LlavaNext, Mistral3, Llama4, etc.) before #1904. However, there are no reported issues or OOMs for those models — likely because their parameter counts are smaller and FSDP alone is sufficient to fit in memory, so even if AC were silently skipped it would not cause an observable failure.

Gemma4 31B is the first model where memory is tight enough that AC is strictly required, which is why this bug surfaced here.

Fix is in #1941 — flattens nn.ModuleList results in _extract_model_layers so individual decoder layers are returned, restoring correct AC for all registered model types.

[athitten]

Thank you so much @khazic for fixing this issue. Appreciate your contribution here!! Fix in #1941 which is merged. Closing the issue

[qiaochuz-nv]

QA RCCA Analysis

Fix PR: #1941 (same fix as #1986)
Root Cause: Gemma4 31B OOM after PR #1904 changed tied lm_head logic in stateful_wrappers.py, which also broke AC layer detection for VLMs.

Trigger Config: Gemma4 31B + VLM SFT + FSDP2

Coverage Assessment:

Dimension	Bug Config	Nearest Test	Match
Gemma4 31B	Required	automodel_gemma4_31b_vlm_finetune_test	Yes
VLM SFT	Required	Same test	Yes
AC	Needed for OOM avoidance	Fixed by #1941	Yes

Status: COVERED — Direct Gemma4 31B VLM test exists. Same root cause as #1986.

Conclusion: No new test needed. Same fix and coverage as #1986.