Fix optimizer CPU offload for megatron-fsdp dtensor param

[wplf]

Summary

• Handle DTensor parameters and gradients by operating on local shards before optimizer CPU offload copies.
• Avoid dispatching pin_memory/is_pinned through DTensor and respect pin_cpu_params.

Tests

• uv run isort megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py
• PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. python -m pytest tests/unit_tests/test_optimizer_cpu_offloading.py -q

test result

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[wplf]

/ok to test

[copy-pr-bot]

/ok to test

@wplf, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[wplf]

/ok to test 5da7d67

[cspades]

@wplf Could you share some information on what the bug was? Curious what was the root cause and how it is related to DTensors!

[cspades]

BTW, if we don't need differentiability (such as in this optimizer step), sometimes dtensor._local_tensor is more robust since we can be sure we do not receive a copy of the original local Tensor. (Don't ask me why I know this.)

[wplf]

Hi, @cspades .

Currently, when using megatron-fsdp and --optimizer-cpu-offload, it will crash. Root cause is hybrid-optimizer does not take mfsdp into consideration. Use Dtensor._local_tensor will resolve this issue.

But I haven't check this fix thoroughly for training convergence.

I'll check training stability and convert Dtensor.to_local() to Dtensor._local_tensor as soon as possible.