[faresobeid]

Ranks are NUMA-bound only with model.offload = "full". With the default model.offload = "optimizer" (and with activation offloading) they are not, so each rank's pinned buffers land on whichever socket its thread happened to run on when it first touched them. On an 8xH200 node (2 sockets, GPUs 0-3 on node 0, 4-7 on node 1), one run put three of eight ranks' optimizer state on the remote socket, and a run with activation offloading put all eight there. Copies to and from the remote socket are slower, and the slow ranks hold back the step for every rank.

Change

• setup_cpu_offload(config.model) replaces setup_full_cpu_optimizer_offload(). It binds the rank to its GPU's NUMA node (the existing bind_process_to_gpu_numa_node) whenever model.offload != "none" or activation offloading is on. It sets the CPU optimizer thread count only for full offload, as before.
• bind_process_to_gpu_numa_node now finds the rank's GPU in NVML by UUID (torch.cuda.get_device_properties(device).uuid) instead of by its CUDA index. NVML numbers every GPU on the node, while CUDA numbers only the visible ones, so under CUDA_VISIBLE_DEVICES remapping (colocated RL) the old lookup bound ranks to another GPU's socket.

Copy bandwidth (pinned host memory, 8xH200, all 8 GPUs at once, GiB/s per GPU)

buffers on	H2D	D2H	both directions at once
GPU's own socket	26.0	25.8	38.8
other socket	26.0	17.6	23.7
unbound (mixed)	25.7	16.0-19.0	22.2-24.4

Trainer (Qwen3-30B-A3B, 1 node 8xH200, FSDP 8 / EP 8, 16k seq, full AC, compile, fake data, mean of steps 6-20)

config	main	this PR
offload = "optimizer", no activation offload	26.8k tok/s, optimizer step 2.33-2.51 s depending on the rank	27.7k tok/s (+3.6%), 2.33 s on every rank
default (offload = "optimizer" + activation offloading)	23.7k tok/s, fwd+bwd 2.49 s, optimizer step 2.83 s	26.4k tok/s (+11.5%), fwd+bwd 2.41 s, optimizer step 2.33 s
with #4008 (offload = "optimizer", no activation offload)	31.8k tok/s, optimizer step 1.28-1.75 s	33.8k tok/s (+6.3%), 1.48 s on every rank

Peak GPU memory and host RSS are unchanged (63.2 GiB, 41-45 GiB per rank). With the binding, each rank's offloaded state sits on its GPU's socket (checked in /proc/self/numa_maps); without it, the socket was effectively random per rank. Loss and grad norm per step stay within the run-to-run noise of main; the change only moves where the host memory lives.

🤖 Generated with Claude Code

---

Note

Low Risk
Startup-only affinity and thread tuning for training; no change to model math, checkpoints, or APIs beyond when binding runs.

Overview
NUMA binding for all host offload paths, not only offload = "full". setup_cpu_offload(config.model) replaces setup_full_cpu_optimizer_offload() and is always invoked from both trainer entrypoints; it binds the process to the GPU’s NUMA node when model.offload != "none" or activation offloading is enabled, and still configures CPU optimizer threads only for full offload.

GPU lookup for binding now goes through NVML by GPU UUID instead of CUDA device index, so ranks stay pinned to the correct socket when CUDA_VISIBLE_DEVICES remaps GPUs (e.g. colocated RL).

Log output for a successful bind includes the PCI sysfs id alongside the CUDA device id.

^{Reviewed by Cursor Bugbot for commit 16dc340. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 8a3d4e2. Configure here.}