[train] FSDP2 cannot train glm4v VLMs: per-nn.Embedding fully_shard leaves vision position_embedding sharded at raw .weight read (grid_sample DTensor crash)

[dzorlu]

Summary

Training any glm4v-family VLM (GLM-4.1V-9B-Thinking, GLM-4.6V-Flash) with trainer.strategy=fsdp crashes in the first policy forward:

File ".../transformers/models/glm4v/modeling_glm4v.py", line 198, in forward
    interpolated_embed_fp32 = F.grid_sample(
...
RuntimeError: aten.grid_sampler_2d.default got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!

Measured on GLM-4.6V-Flash, 4x H200, transformers 5.8.0, GRPO multi-turn with images. Rollout is unaffected (vLLM holds its own weights); the crash is in the FSDP2-wrapped training model.

Mechanism

apply_fsdp2 gives a dedicated fully_shard group to every nn.Embedding when tie_word_embeddings is false:

https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl/backends/skyrl_train/distributed/fsdp_utils.py#L226-L231

A per-module group unshards its parameters only in that module's own forward pre-hook. Word embeddings are always used through their own forward(), so the hook fires and everything works. GLM4V's vision tower breaks the assumption: Glm4vVisionEmbeddings.forward reads self.position_embedding.weight raw (never calling the embedding's forward) and feeds it to F.grid_sample to interpolate position embeddings per image grid. grid_sample has no DTensor sharding rule, the weight is still a sharded DTensor at access time, and the op refuses the mixed inputs.

GLM-4.6V-Flash has top-level tie_word_embeddings: false, so the branch is always armed. No config avoids it: wrap_policy.transformer_layer_cls_to_wrap is OR'd with the embedding branch, and overriding tie_word_embeddings=True via model config kwargs would wrongly tie lm_head at load.

This is inherited from verl, where it is also latent

This selection logic was adapted from verl. verl main has the same per-nn.Embedding wrap in _select_fsdp2_wrap_targets (verl/utils/fsdp_utils.py), but verl does not hit the crash because its default actor strategy is still FSDP1 (verl/trainer/config/actor/dp_actor.yaml: strategy: fsdp, with a TODO(haibin.lin): switch to fsdp2). verl's GLM-4.1V FSDP example runs FSDP1, where forward-time params are plain flat-param views. The same error class has been reported there in another FSDP2 context with no resolution (verl-project/verl#4756). Since SkyRL is FSDP2-only, SkyRL cannot train glm4v at all today.

Fix

Wrap only model.get_input_embeddings() instead of every nn.Embedding. Module selection is identical for every model whose only embedding is the input embedding (all text models; Qwen-VL vision uses a conv patch embed, no auxiliary nn.Embedding), so no behavior change for existing models. Auxiliary embeddings like GLM4V's vision table fall into the root group, whose params are plain unsharded tensors for the whole forward; the memory cost is that table staying unsharded during forward (a few M params for glm4v).

PR incoming.

Validation

• GLM-4.6V-Flash GRPO run on 4x H200 crashes at the first training forward on main; with the fix the forward proceeds.
• Unit tests cover: input embeddings + layer classes selected, auxiliary embedding (GLM pattern) not selected, tied embeddings not selected, models without get_input_embeddings.