[fsdp] Tied embeddings disable rank0-only weight loading, so host RAM scales with rank count; the naive fix unties the model silently

System Info

Read on current main (68c9ac35); reproduced on a9f29851.

• verl/workers/engine/fsdp/transformer_impl.py:262 (_build_module)
• verl/utils/fsdp_utils.py:64 (get_init_weight_context_manager), :505-513 (fsdp2_load_full_state_dict)

Three machines: 8x H200 141GB / 1511GB host RAM / driver 580.173.02; 2x H100 80GB / 503GB host RAM / driver 580.126.09; 4x H100 80GB SXM / 1006GB container memory limit / driver 580.126.09. torch 2.11.0+cu130, transformers 5.9.0, vllm 0.24.0 — the repo's own pins.

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

I ran into this while checking whether cohere2_moe works on FSDP + vLLM GRPO for #7796, using CohereLabs/North-Mini-Code-1.0 (30B/A3B, 61 GB bf16, tie_word_embeddings: true). Two GRPO steps never started: every attempt was OOM-killed during from_pretrained, before a single GPU was touched.

Part 1: tied embeddings disable the rank0-only load

_build_module turns off meta initialisation whenever embeddings are tied:

init_context = get_init_weight_context_manager(
    use_meta_tensor=not self.model_config.hf_config.tie_word_embeddings, mesh=self.device_mesh
)

With use_meta_tensor=False, get_init_weight_context_manager hands every rank cpu_init_weights instead of init_empty_weights, so all N ranks materialise the whole checkpoint on host RAM rather than one rank materialising and the rest receiving weights by broadcast. Host RAM, not VRAM, becomes the bound on single-node loading, and it scales with GPU count.

On the 8x H200 node, with actor.strategy=fsdp2:

peak host RAM   1437 GB / 1511 GB
died in         from_pretrained, at 72% of "Loading weights: 283/394"
GPU             0% utilisation throughout; never reached

Ray's memory monitor killed a worker first (Killing 1 worker(s) ... 1336.05GB / 1405.37GB (0.950678)); with RAY_memory_monitor_refresh_ms=0 it reached 72% instead of 64% and then a worker died abruptly (ActorUnavailableError ... RpcError: Socket closed) with host RAM at the 1437 GB peak above. dmesg is not readable in the container so I cannot show an OOM-killer line for that second death. param_offload=False / optimizer_offload=False made no difference, which is what pointed at load time rather than training state.

For honesty about the magnitude: 8 x 61 GB = 488 GB alone would have fit in 1511 GB. This checkpoint stores its experts unfused — 18432 per-expert tensors that HF fuses into 394 parameters — so each rank also holds fusion buffers (source and destination together) on top of its own copy, which is what pushed it over. Measured per rank in the follow-up below, the load costs ~112 GB (VmHWM) rather than 61 GB — the ~180 GB I first derived by dividing the 8-rank peak had folded in the rest of the GRPO stack. The guard is what multiplies that per-rank cost by N.

Part 2: the obvious fix unties the model, silently

FSDP2 already has the machinery to avoid all of this: fsdp2_load_full_state_dict materialises non-zero ranks with to_empty() and fills them via set_model_state_dict(..., broadcast_from_rank0=True). So the guard reads like an FSDP1-era leftover, and #5746 proposed exactly that — take meta init unconditionally on fsdp2. Its author self-closed it 35 hours later with no human review and no stated reason, and has not been back; nothing has touched these lines since.

That fix is not sufficient on its own. to_empty() gives every parameter fresh storage, so a tied lm_head stops aliasing the input embedding:

meta init:              lm_head is embed_tokens -> True
after to_empty('cpu'):  lm_head is embed_tokens -> False
after tie_weights():    lm_head is embed_tokens -> True

Nothing fails when that happens. set_model_state_dict fills both entries from rank 0's state dict — which does carry lm_head.weight as an alias, so the values are right — the model runs, and the loss falls. But gradients no longer reach the shared tensor from both paths, so the embedding receives only the input-side half of its gradient:

embed_tokens.grad norm, tied   (rank0 path)   : 0.354796
embed_tokens.grad norm, untied (to_empty path): 0.227071

The broadcast ranks would train a different model from rank 0, with no error and no warning. This is presumably why the guard exists at all: it trades wasted host RAM for not being wrong.

Scope

Deliberately not overstating this. The trigger is a model that is both tied and large, and in practice tying is a small-model feature — the large models verl users commonly train are untied:

So for most users the redundant load is a few GB and invisible. It bites on the unusual large-and-tied checkpoints, of which North-Mini-Code-1.0 is one. It also does not look like a verl bug from the outside — it presents as "OOM while loading a big model", which is easy to attribute to the machine and work around by using fewer GPUs.

Worth noting that #5746's author reported a different symptom from the same guard — use_meta_tensor=False leaving FSDP2 sharded-param references on CPU and causing gradient device mismatches in backward. I did not reproduce that one, but two independent symptoms from one line is part of why I think it is worth fixing rather than documenting.

Expected behavior

A tied-embedding model on fsdp2 should load the way an untied one does: rank 0 materialises, the other ranks build on meta and receive weights by broadcast, and the tie is intact on every rank afterwards.

That needs both halves — the strategy check #5746 proposed, and a re-tie after to_empty():

# transformer_impl.py
use_meta_tensor = self.engine_config.strategy == "fsdp2" or not tie_word_embeddings

# fsdp_utils.py, after to_empty() in fsdp2_load_full_state_dict
model.tie_weights()

I have this working and will open a PR referencing this issue. What it carries that #5746 did not:

• CPU regression tests pinning that to_empty drops the tie, that tie_weights() restores it, and that the embedding gradient matches the rank-0 module only when re-tied (the un-retied case is asserted to differ, so the test cannot pass vacuously).
• GPU validation on 2x H100 with Qwen2.5-0.5B (tie_word_embeddings: true): rank 1 builds on meta while rank 0 materialises; lm_head.weight is embed_tokens.weight holds on both ranks after fully_shard + broadcast; the SFT engine trains end to end (train/loss 0.912957, val/loss 0.742465).
• A regression check that the untouched FSDP1 path is unaffected — and on the same model its first-step loss is 0.912957489490509, bit-identical to the patched fsdp2 path, which is the strongest evidence I have that the tie really survives.

Happy to split the two halves, or to drop the strategy check and only fix the silent untying, if maintainers would rather keep the guard and treat the memory behaviour as intended.

Follow-up: measured on the originating checkpoint with the fix

Same model, 4x H100 80GB, 2 TB host RAM (1006 GB container limit), a9f29851 + the fix in #7835. Load-only through _build_module's exact sequence, each rank's VmHWM plus the host's free used-memory peak:

                       rank0      rank1     rank2     rank3    host "used" peak
use_meta_tensor=False  112.0 GB   112.0 GB  112.0 GB  112.0 GB   318 GB   (today's behaviour)
use_meta_tensor=True   112.0 GB    56.7 GB   56.7 GB   56.7 GB   143 GB   (rank 0 only; ranks 1-3 on meta)
tie held on every rank:  True in both runs

The ~57 GB on the meta ranks is file-backed (from_pretrained still maps the safetensors and RSS counts touched file pages); the host-level used figure, which excludes page cache, is the one that reflects anonymous memory. So on this checkpoint the fix takes the load from N x ~112 GB of anonymous memory to one rank's worth, with lm_head still aliasing the embedding on every rank afterwards.

What I could not show: a completed training step on the 30B on 4x 80 GB. Three SFT attempts with the fix stalled (details in #7835). The control I ran to attribute that — unpatched a9f29851, same engine.offload_policy=True command, identical 4x H100 pod with a 1006 GB container memory limit — instead reproduced this issue through the SFT trainer: from_pretrained completed on all four ranks (fp32 master weights, so ~4x the bf16 footprint), host memory reached the cgroup limit (memory.peak == memory.max, oom_kill 1), and rank 2 was SIGKILLed about 3.5 minutes in, before any parameter shard reached a GPU. The patched run on the same pod type loads on rank 0 only and gets through fully_shard at 16-20 GB per GPU. So the memory half of this issue is confirmed on 4 GPUs with the real trainer; the stall after loading remains unattributed because unpatched code cannot reach that point on this hardware.

Update: on Qwen2.5-0.5B (tied), 2x H100, the fix under engine.offload_policy=True trains with step-1 loss bit-identical to the unpatched path and step-2 losses within run-to-run noise (table in #7835), so the 30B stall does not reproduce at small scale.