[Orange-Lv]

Summary

This PR preserves pinned host memory for CPU-offloaded FP32 params in HybridDeviceOptimizer.

In the optimizer_cpu_offload=True + param_update_in_fp32=True path, the offloaded CPU param is first created as pinned host memory, but a later clone().float() creates a new pageable CPU tensor. As a result, the main HtoD copy-back path can fall back to pageable host memory.

This change re-applies pin_memory() to the final FP32 CPU param used for GPU copy-back when pin_cpu_params=True.

Problem

In megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py, the offloaded CPU param is created like this:

if offload_params_numel < offload_threshold and param.is_cuda:
    param = param.detach().clone().cpu().pin_memory()
    offload_params_numel += param.numel()
    cpu_copy = True

But in the FP32 update path, it is later replaced with:

if self.param_update_in_fp32 and param.dtype != torch.float32:
    param = param.detach().clone().float()
    param_to_fp32_param[orig_param] = param

The second clone().float() produces a new CPU tensor and does not preserve the original pinned-memory property.

This matters because HybridDeviceOptimizer does not only offload optimizer state. For CPU-side optimizer execution, it also:

• copies gradients from GPU to CPU,
• updates CPU-side parameter copies on CPU,
• copies updated parameters back to GPU.

That final copy-back happens here:

gpu_param.data.copy_(param.data, non_blocking=True)

If param.data is pageable instead of pinned, the HtoD path is slower and may show up as pageable host memory in traces.

Change

Re-apply pin_memory() after the FP32 clone for CPU-offloaded params when pin_cpu_params=True:

if self.param_update_in_fp32 and param.dtype != torch.float32:
    param = param.detach().clone().float()
    if cpu_copy and self.pin_cpu_params:
        param = param.pin_memory()
    param_to_fp32_param[orig_param] = param

This keeps the change minimal and only affects the CPU-offload + FP32 path.

Validation

1. Functional correctness

Ran a short GPT2-1.5B mock-data training with:

• Transformer Engine + Apex
• optimizer_cpu_offload=True
• optimizer_offload_fraction=1.0

Smoke test:

• train-iters=20
• completed successfully
• no new runtime errors
• loss/step behavior remained normal

Observed log:

• iter 1: 17858.3 ms, loss 1.113179E+01
• iter 10: 5016.5 ms, loss 1.063229E+01
• iter 20: 3606.6 ms, loss 9.541001E+00

2. Trace evidence

Compared nsys traces before and after the patch.

Before:

• dominant HtoD path:
  • (HTOD, Pageable -> Device): 20374 calls
• pinned HtoD path:
  • (HTOD, Pinned -> Device): 1360 calls

After:

• dominant HtoD path:
  • (HTOD, Pinned -> Device): 6200 calls
• pageable HtoD path:
  • (HTOD, Pageable -> Device): 513 calls

So the main HtoD copy-back path changed from Pageable -> Device to Pinned -> Device.

Average duration for the dominant HtoD path:

• before: 2,362,300 ns (~2.36 ms)
• after: 819,103 ns (~0.82 ms)

This indicates the patch not only changes the memory kind shown in traces, but also improves the HtoD copy-back path.

Notes

The before/after nsys captures were not identical in total duration, so aggregate percentages should be interpreted qualitatively.

However, the key result is robust:

• the dominant HtoD path changes from pageable to pinned,
• and its average copy duration decreases significantly.

Why this is safe

This change does not expand pinning globally. It only preserves pinning for the final FP32 CPU tensor used in GPU copy-back, and only when:

• the param is actually CPU-offloaded, and
• pin_cpu_params=True.

So this aligns runtime behavior with the existing user-facing configuration rather than introducing a broader semantic change.

[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.

[dimapihtar]

LGTM. Thank you!

[dimapihtar]

/ok to test 7687e8e

[guihong-nv]

/ ok to test fa99fc9

[guihong-nv]

/ok to test 517c792

[guihong-nv]

/ok to test 83819f1

[guihong-nv]

/ok to test 845e8c1

[ilml]

/ok to test 845e8c1

[janEbert]

Dear @Orange-Lv (and others, too, I believe? Not sure how merge commits are handled WRT this), could you add a DCO to your commits using git commit -s? We recently started to require all commits to be signed with DCOs and that's currently preventing the merge of this PR.

[maanug-nv]

/ok to test 4b34420

[maanug-nv]

/ok to test 2b8af82

[maanug-nv]

/ok to test f6c5eff

[maanug-nv]

/ok to test 5f73bc8

[nemo-automation-bot]

🔄 Merge queue validation started!

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