Hi @yuhezhang-ai , does this handle multiple compute dtypes in the frozen part? Also i see that it iterates over hte model/model.parts, so if this explicitly processes the frozen part or it iterates over the whole model and skips the wrapped parts. I feel explicit is better, because we can't assume hard FSDP2 as the downstream wrapper.

Thanks. Based my latest code:

It casts a frozen subtree to a single compute dtype (FSDP's param_dtype), with fp32 carve-outs for _keep_in_fp32_modules(_strict).

Iterate-and-skip vs. explicit / FSDP2 assumption: fair point. The pass walks the whole model and casts each fully-frozen submodule — buffers always, plain params always. The one assumption is the DTensor branch: a sharded param is skipped on the theory that the wrapper re-gathers and down-casts it to the compute dtype, which is true for FSDP2 but not guaranteed for every DTensor-producing wrapper. In practice it's scoped to FSDP today because the whole pass only runs when there's an FSDP mp_policy.param_dtype, so other wrappers don't hit this path.
But I agree explicit is cleaner: the robust fix is to handle the frozen tower explicitly (exclude it from sharding so it stays plain and gets cast directly), which may also help perf. I think that's out of scope for this PR and gemma4 may need a small redesign — see my later comments on gemma4.

Hi @yuhezhang-ai have you confirmed this working?

Hi @akoumpa, thanks!

On gemma4 specifically: it was indeed problematic. I dug into this. The frozen vision tower is actually sharded (it lands in the root FSDP unit, since gemma4 runs with wrap_outer_model=True), so FSDP down-casts its params to bf16 at all-gather. The actual leak is its buffers: HF's Gemma4VisionModel keeps fp32 std_bias/std_scale buffers and FSDP mixed precision never casts buffers, so (hidden_states - std_bias) * std_scale promotes the bf16 activation back to fp32 right before the bf16 multimodal projector. That's the mismatch. The fix 0915fd1 casts frozen buffers to the compute dtype, which removes it (verified the rest of the path — pooler, projector norm — all cast back).
Now it's running correctly after my fix: https://wandb.ai/Nemo-automodel/yuhez_workspace/groups/gemma4-26b-a4b-dtype-verify-rerun/workspace?nw=nwuseryuhez

Follow-up (separate PR) if we want to exclude vision_tower from fsdp sharding : two things combine here, gemma4 runs with wrap_outer_model=True, and the MoE parallelizer's frozen-tower skip (moe/parallelizer.py apply_fsdp()- around line 331) only checks audio_tower/visual, not vision_tower (which is also nested at model.model.vision_tower). So the frozen tower slips past the skip and the outer wrap then sweeps it into the root unit (sharded). To truly exclude it we'd need to fix both: add vision_tower to the skip and reconsider wrap_outer_model here (I'm not sure flipping it is safe / what else it affects). That's really a question for who owns the gemma4 model design. For now the buffer cast fixes correctness without touching the sharding layout.