Avoid extra MFSDP v2 model-weight sync memcpy

[wujingyue]

Summary

• Avoid an extra device-to-device memcpy when MFSDP v2 syncs optimizer weights into compute weights and the placements already match.
• Cast main_weight directly into the persistent model_weight buffer for the same-placement path.
• Keep cast-before-redistribute for different placements so cross-rank communication uses the lower-precision compute dtype instead of the higher-precision optimizer dtype.

Root cause

sync_model_weight_from_main_weight() always cast main_weight into a temporary compute-dtype DBuffer, then called redistribute(..., out=model_weight). When main_weight and model_weight already had the same placements, redistribute only copied the temporary buffer into model_weight, producing an avoidable D2D memcpy after the cast.

Implementation

This adds an out= path to DBuffer.cast and uses it from MFSDP v2 model-weight sync. The DBuffer helper change is an implementation detail for writing the cast result directly into existing storage.

Validation

• UV_NO_SYNC=1 CUDA_VISIBLE_DEVICES=0,1 uv run python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py
  • 19 passed, 6 skipped
• UV_NO_SYNC=1 CUDA_VISIBLE_DEVICES=0,1 uv run python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
  • 14 passed
• git diff --check
  • passed

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[wujingyue]

/ok to test 85cbc0c

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29602452616