Repository navigation
[MFSDP V2] Post-wrap load_state_dict(assign=True) does not resume from the loaded checkpoint #6087
Description
Observed behavior
At bc67abde107ca0a7751208ae142d09ce602e45be, post-wrap
load_state_dict(checkpoint, assign=True) returns a clean load. For the fixed
identity input and zero target, the pre-load weights have loss 0.390625,
while the checkpoint weights have loss 0.44921875. The measured loss after
the load is 0.390625000, which matches the pre-load weights rather than the
checkpoint weights.
Verified with PyTorch 2.8.0+cu128, CUDA 12.8, and two NVIDIA
A100-SXM4-80GB GPUs. The complete stdout from each of two fresh DP=2 runs was:
NCCL version 2.27.3+cuda12.9
load_result: missing_keys=[], unexpected_keys=[]
step | loss
0 | 0.390625000
1 | 0.390625000
2 | 0.390625000
3 | 0.390625000
4 | 0.390625000
5 | 0.390625000
6 | 0.390625000
7 | 0.390625000
8 | 0.390625000
9 | 0.390625000
10 | 0.390625000
11 | 0.390625000
12 | 0.390625000
13 | 0.390625000
14 | 0.390625000
15 | 0.390625000
16 | 0.390625000
17 | 0.390625000
18 | 0.390625000
19 | 0.390625000
20 | 0.390625000
loss_decreased=False
This is a deterministic convex MSE problem: a bias-free linear layer receives
the identity matrix as input and is trained toward a zero target with SGD. The
checkpoint contains fixed nonzero weights and comes from the normal
state_dict() of an identically sharded model, with every value frozen by
detach().clone(). The optimizer is constructed only after the assign=True
load. The measured loss is identical at every one of the 20 steps: the model does not train, and it stays at the loss of its own pre-load weights.
Both fresh torchrun processes exited 0 and produced byte-identical stdout.
Minimal reproduction
Save the following as repro_issue_6087_assign_true_only.py in the
Megatron-LM repository root:
#!/usr/bin/env python3
"""Behavior-only reproduction for post-wrap load_state_dict(assign=True)."""
from __future__ import annotations
import os
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import nn
from torch.distributed.device_mesh import init_device_mesh
from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import (
Flat,
Placements,
fully_shard,
)
WORLD_SIZE = 2
SEED = 6087
STEPS = 20
LEARNING_RATE = 0.5
def build_model() -> nn.Linear:
return nn.Linear(4, 4, bias=False)
def placements() -> Placements:
return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()])
def set_weight(model: nn.Linear, values: list[list[float]]) -> None:
with torch.no_grad():
model.weight.copy_(torch.tensor(values, dtype=torch.float32, device=model.weight.device))
def evaluate_loss(model: nn.Module, inputs: torch.Tensor, targets: torch.Tensor) -> float:
with torch.no_grad():
loss = F.mse_loss(model(inputs), targets)
dist.all_reduce(loss, op=dist.ReduceOp.SUM)
return (loss / dist.get_world_size()).item()
def main() -> None:
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl")
try:
rank = dist.get_rank()
world_size = dist.get_world_size()
if world_size != WORLD_SIZE:
raise RuntimeError(f"Expected world_size={WORLD_SIZE}, got {world_size}")
torch.manual_seed(SEED)
torch.use_deterministic_algorithms(True)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
device = torch.device("cuda", local_rank)
mesh = init_device_mesh("cuda", (world_size,))
shard_placements = placements()
inputs = torch.eye(4, dtype=torch.float32, device=device)
targets = torch.zeros(4, 4, dtype=torch.float32, device=device)
source_model = build_model().to(device)
set_weight(
source_model,
[
[1.00, -0.50, 0.25, -0.75],
[-0.25, 0.75, -1.00, 0.50],
[0.50, 0.25, -0.50, 1.00],
[-0.75, 1.00, 0.50, -0.25],
],
)
fully_shard(source_model, mesh=mesh, placements=shard_placements)
checkpoint = {
name: value.detach().clone() for name, value in source_model.state_dict().items()
}
model = build_model().to(device)
set_weight(
model,
[
[0.25, -0.75, 0.50, -0.25],
[0.50, 0.25, -0.75, 1.00],
[-1.00, 0.50, 0.25, -0.50],
[0.75, -0.25, 1.00, 0.50],
],
)
fully_shard(model, mesh=mesh, placements=shard_placements)
load_result = model.load_state_dict(checkpoint, assign=True)
optimizer = torch.optim.SGD(
model.parameters(),
lr=LEARNING_RATE,
momentum=0.0,
weight_decay=0.0,
foreach=False,
)
losses = [evaluate_loss(model, inputs, targets)]
for _ in range(STEPS):
optimizer.zero_grad(set_to_none=True)
loss = F.mse_loss(model(inputs), targets)
loss.backward()
optimizer.step()
losses.append(evaluate_loss(model, inputs, targets))
loss_decreased = losses[-1] < losses[0]
if rank == 0:
print(
"load_result: "
f"missing_keys={list(load_result.missing_keys)}, "
f"unexpected_keys={list(load_result.unexpected_keys)}"
)
print("step | loss")
for step, loss in enumerate(losses):
print(f"{step:>4} | {loss:.9f}")
print(f"loss_decreased={loss_decreased}", flush=True)
dist.barrier()
finally:
if dist.is_initialized():
dist.destroy_process_group()
if __name__ == "__main__":
main()Run from the repository root:
# CUBLAS_WORKSPACE_CONFIG is required by torch.use_deterministic_algorithms(True).
CUBLAS_WORKSPACE_CONFIG=:4096:8 NCCL_DEBUG=WARN \
torchrun --standalone --nproc-per-node=2 \
repro_issue_6087_assign_true_only.pyExpected behavior
Post-wrap load_state_dict(assign=True) should make continuation training use
the loaded parameters, or it should be explicitly rejected during load with a
clear error. It should not return clean success and then leave the model neither loaded nor training.
Activity
chengcuiping commentedon Jul 28, 2026
@wujingyue — following up on the design document requested in your review of #5836.
My proposed direction is V2-first: keep pre-wrap loading as the default and, if post-wrap replacement is an MFSDP-supported contract, initially limit support to a pre-execution state-dict path through holder-owned reset or reconciliation. Post-optimizer replacement would remain unsupported.
The primary decision is whether post-wrap
Parameterreplacement should be supported by MFSDP at all or treated as a caller error. I’ll reduce, redirect, or close #5836 based on that answer.If a regression-first step would help, I can turn the existing two-rank V2 DP probe into a focused test with a non-uniform value oracle before proposing any API.
wujingyue commentedon Jul 28, 2026
Thanks for putting this together.
Could you rewrite the doc to focus on V2 only?
Also, could you make the problem statements more concrete with minimal reproductions? I'd lean toward practical examples over philosophical arguments, so it's easier to see the motivation for each design decision.
chengcuiping commentedon Jul 29, 2026
Thanks — that makes sense.
I’ll rewrite the document to focus exclusively on MFSDP V2 and replace the broad lifecycle discussion with concrete, runnable examples. I’ll first rerun the existing DP=2
load_state_dict(assign=True)case on currentmain, using a minimal model and non-uniform values, and include pre-wrap loading andassign=Falseas controls.The revised document will include the exact commands and observed outputs through the normal forward path, then connect each design decision directly to those results. I’ll remove the V1 and general
reregister_parameters()discussion, keep #5836 paused, and update the issue body once the reproductions are verified.chengcuiping commentedon Jul 29, 2026
Thanks — I rewrote the issue around MFSDP V2 and included the complete 181-line DP=2 reproduction inline.
At
bc67abde107ca0a7751208ae142d09ce602e45be, the exact inline reproduction was run twice independently under DP=2. Both runs produced identical results. In A and C, both ranks retained the exact source shards and ordinary forward matched the source reference. In B, post-wrapassign=Truereturned a clean load and installed the source value on the module while the cached holders retained the target value; ordinary forward then used and restored the old target holders.No production V2 caller was identified within the recorded search scope. I therefore recommend rejecting post-wrap
assign=Trueat the state-dict load boundary, before module or holder state mutates, with pre-wrap loading and the demonstrated same-layoutassign=Falsepath as the actionable alternatives. A narrow holder reset remains an option only if a concrete caller establishes the need.Could you confirm whether this unsupported-at-load contract is desired? #5836 remains Draft and paused pending alignment.
wujingyue commentedon Jul 29, 2026
Since
load_state_dict(assign=True)is the only concrete problem, can you rewrite this issue to just focus on that? Also, can you post/inline the repro?[-][Design Doc] Parameter replacement semantics for Megatron-FSDP[/-][+][MFSDP V2] Post-wrap load_state_dict(assign=True) does not affect the next forward[/+]chengcuiping commentedon Jul 30, 2026
Thanks @wujingyue — and sorry: in the previous revision, the reproduction was folded inside a
Details
block, so it read as missing.I rewrote the issue to focus only on experimental MFSDP V2 post-wrap load_state_dict(assign=True). The complete reproduction and actual rank-by-rank output are now inline in the body. At bc67abd, I reran the exact script twice under DP=2; both ranks completed all exact assertions, and the two captured stdout files were byte-identical.
The call currently returns success, but the loaded value does not survive the first ordinary forward. I recommend rejecting this sequence during load, before module or holder state mutates. Would that be the desired MFSDP V2 contract?
wujingyue commentedon Jul 30, 2026
The repro as is seems to add quite some artificial assertions around underlying mechanisms. Can you make it more behavior focusing? For example, we could build a model using TE / MCore / plain PyTorch without artificial attributes and then show the training doesn't converge compared to some baseline?
[-][MFSDP V2] Post-wrap load_state_dict(assign=True) does not affect the next forward[/-][+][MFSDP V2] Post-wrap load_state_dict(assign=True) does not resume from the loaded checkpoint[/+]chengcuiping commentedon Jul 31, 2026
Thanks, @wujingyue — I rewrote the reproduction around a standard
nn.Sequentialmodel and a pure behavior-level A/B/C comparison. A and B usethe same cloned mapping from
source_model.state_dict(); C does not load.In two fresh DP=2 runs, A and B both reported clean loads. At the first forward,
source vs A max diff was 0.000000000000000, A vs B was 0.812111496925354, and B
vs no-load C was 0.000000000000000. Across the fixed 20-step continuation, the
losses were A: 0.018781833 -> 0.018759176, B: 0.193090320 -> 0.193090320, and C:
0.193090320 -> 0.062155873. All three optimizers are constructed at the same
point, after their load/no-load operation.
The two captured stdout files were byte-for-byte identical (SHA-256:
437350b5db4d26114c60671ca06bd47e2e6292dedb985035100766ec2ac37ca1).The updated issue body contains the complete script, public command, and
verbatim output.
Would you prefer MFSDP V2 to support post-wrap
assign=Truecheckpoint loads,or explicitly reject this combination with a clear error? I can prepare the
corresponding focused patch.
5 remaining items
wujingyue commentedon Aug 1, 2026
Simplify the repro to focus on assign=True case. A simple check for that is to see if loss is decreasing at all during continued training, so we don't really need a baseline.
chengcuiping commentedon Aug 1, 2026
Thanks @wujingyue — reduced to the
assign=Truecase only, no baseline.A bias-free
nn.Linear(4, 4)gets the identity matrix as input and a zero target, so the loss is just the mean square of the weights. The pre-load weights give0.390625; the checkpoint weights would give0.44921875.The post-wrap load returns clean, and the measured loss is
0.390625at every step from 0 through 20 — the model stays on its pre-load weights and does not train.Body updated with the short script and complete output; both DP=2 runs exited 0 with identical stdout.
wujingyue commentedon Aug 1, 2026
Repro looks much better!
Now explore options. Try to be concise. Don't write PRs yet.
chengcuiping commentedon Aug 2, 2026
Short version:
assign=Truereplaces the module-visibleParameter, while MFSDP's next forward continues using its existing holder-owned state. I explored four options; no PR has been prepared.FsdpModulereplaces its parameters.assign_to_params_buffersload metadata. Recursive loading is not a whole-model transaction, so earlier siblings may already have loaded before the guard is reached.0.390625, rather than the checkpoint value0.44921875, with no decrease over 20 steps.assign=Falsesemantics. True replacement reaches hooks, aliases, mixed precision, and optimizer-facing parameter identities, so it is not locally bounded.assign=Truereturning success while forward and training use different state.Given these results, I recommend A. The error can direct callers to load before
fully_shard()or use an in-place load path (assign=False, or the DCP path in #6024).#6024 is unaffected: it loads into existing sharded tensors and calls
set_model_state_dictwith default options, which useassign=False.Does this load-time rejection contract sound right?
wujingyue commentedon Aug 2, 2026
Option A sounds good. Open a PR.