Skip to content

[bug] Quickstart Qwen3.6 SFT crash with metadata shape mismatch #90

Description

@bogdansalyp

Problem

Quick start sft_qwen3_6_35b.sh crashes with metadata shape mismatch error.

Refers to pytorch/pytorch#166926

Minimal repro

MODEL_PATH=Qwen/Qwen3.6-35B-A3B bash examples/scripts/quick_start/sft_qwen3_6_35b.sh

Expected behavior

No crash

Affected area

area:llm

Environment

docker_run.sh on latest main
B200
CUDA 13.2

Logs

[rank5]: Traceback (most recent call last):
[rank5]:   File "<frozen runpy>", line 198, in _run_module_as_main
[rank5]:   File "<frozen runpy>", line 88, in _run_code
[rank5]:   File "/molt/molt/cli/train_sft.py", line 282, in <module>
[rank5]:     train(args)
[rank5]:   File "/molt/molt/cli/train_sft.py", line 143, in train
[rank5]:     trainer.fit(args, consumed_samples, num_update_steps_per_epoch)
[rank5]:   File "/molt/molt/trainer/sft_trainer.py", line 269, in fit
[rank5]:     logs_dict, window_frac = self._run_microbatch(
[rank5]:                              ^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/molt/molt/trainer/sft_trainer.py", line 182, in _run_microbatch
[rank5]:     self.strategy.backward(
[rank5]:   File "/molt/molt/trainer/fsdp/strategy.py", line 431, in backward
[rank5]:     loss.backward()
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/_tensor.py", line 623, in backward
[rank5]:     torch.autograd.backward(
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/autograd/__init__.py", line 395, in backward
[rank5]:     _engine_run_backward(
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/autograd/graph.py", line 979, in _engine_run_backward
[rank5]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/checkpoint.py", line 1187, in unpack_hook
[rank5]:     frame.check_recomputed_tensors_match(gid)
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/checkpoint.py", line 918, in check_recomputed_tensors_match
[rank5]:     raise CheckpointError(
[rank5]: torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
[rank5]: tensor at position 72:
[rank5]: saved metadata: {'shape': torch.Size([1296, 2048]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=5)}
[rank5]: recomputed metadata: {'shape': torch.Size([1297, 2048]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=5)}
[rank5]: tensor at position 73:
[rank5]: saved metadata: {'shape': torch.Size([1296, 1024]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=5)}
[rank5]: recomputed metadata: {'shape': torch.Size([1297, 1024]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=5)}
[rank5]: tensor at position 74:
[rank5]: saved metadata: {'shape': torch.Size([1296, 1]), 'dtype': torch.float32, 'device': device(type='cuda', index=5)}
[rank5]: recomputed metadata: {'shape': torch.Size([1297, 1]), 'dtype': torch.float32, 'device': device(type='cuda', index=5)}
[rank5]: tensor at position 77:
[rank5]: saved metadata: {'shape': torch.Size([1296, 512]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=5)}
[rank5]: recomputed metadata: {'shape': torch.Size([1297, 512]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=5)}
[rank5]: .

[rank5]: Tip: To see a more detailed error message, either pass `debug=True` to
[rank5]: `torch.utils.checkpoint.checkpoint(...)` or wrap the code block
[rank5]: with `with torch.utils.checkpoint.set_checkpoint_debug_enabled(True):` to
[rank5]: enable checkpoint‑debug mode globally.

[rank5]: If this error occurs under torch.compile with automatic_dynamic_shapes enabled,
[rank5]: it may be because the recomputation selected a different compiled graph than the
[rank5]: forward pass (e.g., a dynamic graph instead of the original static graph).
[rank5]: To fix this, either:
[rank5]:   - Use torch._dynamo.mark_dynamic() to explicitly mark varying dimensions as
[rank5]:     dynamic upfront, avoiding the static-to-dynamic transition.
[rank5]:   - Call torch._C._dynamo.eval_frame._set_lru_cache(False) to disable LRU cache
[rank5]:     reordering, which can change which graph is checked first between forward
[rank5]:     and recompute.
[rank5]: See https://github.com/pytorch/pytorch/issues/166926 for more details and
[rank5]: workaround examples.

Train epoch:   0%|                                                                                                                      | 0/1 [01:06<?, ?it/s]
Train step of epoch 0:   0%|                                                                                                          | 0/262 [01:06<?, ?it/s]
W0820 23:45:59.508000 32 torch/distributed/elastic/multiprocessing/api.py:1028] Sending process 101 closing signal SIGTERM
W0820 23:45:59.509000 32 torch/distributed/elastic/multiprocessing/api.py:1028] Sending process 102 closing signal SIGTERM
W0820 23:45:59.509000 32 torch/distributed/elastic/multiprocessing/api.py:1028] Sending process 103 closing signal SIGTERM
W0820 23:45:59.509000 32 torch/distributed/elastic/multiprocessing/api.py:1028] Sending process 106 closing signal SIGTERM
W0820 23:45:59.509000 32 torch/distributed/elastic/multiprocessing/api.py:1028] Sending process 107 closing signal SIGTERM
E0820 23:46:00.637000 32 torch/distributed/elastic/multiprocessing/api.py:1002] failed (exitcode: 1) local_rank: 3 (pid: 104) of binary: /usr/bin/python3
Traceback (most recent call last):
  File "/usr/local/bin/torchrun", line 6, in <module>
    sys.exit(main())
             ^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 1028, in main
    run(args)
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 1019, in run
    elastic_launch(
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 194, in __call__
    return launch_agent(
           ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 383, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
molt.cli.train_sft FAILED
------------------------------------------------------------

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions