[train] FSDP2: wrap only input embeddings, not every nn.Embedding (fixes glm4v)

[dzorlu]

Fixes #2087.

What

apply_fsdp2 wrapped every untied nn.Embedding in its own fully_shard group. This PR narrows that to model.get_input_embeddings() only, via a new modules_to_wrap_fsdp2() selection function, plus unit tests.

Why

A per-module group unshards its params only in that module's own forward pre-hook. GLM4V's vision position_embedding weight is read raw from the parent forward and fed to F.grid_sample (no DTensor rule), so any glm4v model crashes at the first training forward:

RuntimeError: aten.grid_sampler_2d.default got mixed torch.Tensor and DTensor

Measured on GLM-4.6V-Flash (4x H200, transformers 5.8.0, GRPO multi-turn); details and mechanism in #2087.

Behavior change

None for existing models: for any model whose only nn.Embedding is the input embedding (all text models; Qwen-VL vision towers use conv patch embeds), the selected module set is byte-identical. Auxiliary embeddings (the GLM4V vision table) now land in the root group — plain unsharded tensors during forward, which is the fix; the cost is that table staying unsharded through the forward (~1M params for glm4v).

Note this logic was adapted from verl, and verl main carries the same latent bug in _select_fsdp2_wrap_targets; verl avoids it only because its default strategy is still FSDP1 (dp_actor.yaml has TODO(haibin.lin): switch to fsdp2). Same fix applies there.

Tests

tests/backends/skyrl_train/distributed/test_fsdp2_wrap_selection.py:

• input embeddings + configured layer classes selected
• auxiliary nn.Embedding inside a vision-embeddings-style module NOT selected (the crash case)
• tied embeddings not selected
• model without get_input_embeddings handled

With the fix, a GLM-4.6V-Flash GRPO run proceeds past the training forward that crashes on main.

---

Note

Medium Risk
Changes FSDP2 wrap selection, which affects how parameters are sharded during training. Scope is small and behavior is intended to be identical for models whose only embedding is the input table.

Overview
Fixes GLM-4V FSDP2 training crashing with mixed Tensor/DTensor on F.grid_sample by no longer giving every nn.Embedding its own fully_shard group.

apply_fsdp2 now selects wrap targets via modules_to_wrap_fsdp2: transformer layers plus get_input_embeddings() when embeddings are untied. Auxiliary embeddings (e.g. GLM4V vision position_embedding, whose weight is read from a parent forward) stay in the root group so they are unsharded during forward. Typical text / Qwen-VL wrap sets are unchanged.

Adds unit tests covering input vs vision embeddings, tied embeddings, and models without get_input_embeddings.

^{Reviewed by Cursor Bugbot for commit 0e4e874. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request refactors the FSDP2 wrapping logic to ensure only input (word) embeddings and transformer layers get dedicated shard groups, rather than wrapping all nn.Embedding modules. This specifically prevents crashes in models like GLM-4.6V-Flash where vision position embeddings were incorrectly wrapped. Additionally, a new test suite has been added to verify this behavior. The review feedback suggests making the configuration attribute access safer by using getattr to prevent potential AttributeErrors when model.config or tie_word_embeddings is missing.

[gemini-code-assist]

Directly accessing model.config.tie_word_embeddings can raise an AttributeError if the model does not have a config attribute, or if the config does not define tie_word_embeddings. To make this utility function more robust and support custom or non-standard models (including simpler test models), use getattr to safely retrieve these attributes with sensible defaults.

Suggested change

	input_embeddings = model.get_input_embeddings() if hasattr(model, "get_input_embeddings") else None
	modules = []
	for name, module in model.named_modules():
	if module.__class__.__name__ in transformer_layer_cls_to_wrap or (
	module is input_embeddings and not model.config.tie_word_embeddings
	):
	modules.append(module)
	input_embeddings = model.get_input_embeddings() if hasattr(model, "get_input_embeddings") else None
	config = getattr(model, "config", None)
	tie_word_embeddings = getattr(config, "tie_word_embeddings", False) if config is not None else False
	modules = []
	for name, module in model.named_modules():
	if module.__class__.__name__ in transformer_layer_cls_to_wrap or (
	module is input_embeddings and not tie_word_embeddings
	):
	modules.append(module)