Make CPU optim multithreading by overwriting num-thread but without o…

skills/training/start-run/SKILL.md

@@ -52,6 +52,12 @@ selected environment before starting the run.
   when the defaults are insufficient. It does not create a Gloo process group.
   Set `cpu_optimizer_backend = "torch"` inside the offload table to use fused
   PyTorch AdamW for debugging or parity checks.
+- Any `optim_cpu_offload` mode raises the trainer's intra-op thread count at
+  startup. Launchers export `OMP_NUM_THREADS=1`, which would otherwise leave the
+  bandwidth-bound CPU AdamW kernels on a single core. Each rank claims
+  `cpu_count / local_world_size` threads, capped by its affinity mask, so the
+  ranks on a node never oversubscribe it. This overrides `OMP_NUM_THREADS`, so
+  setting it in `env_vars` does not change the offloaded optimizer's thread count.
 
 ## `rl` — RL training
 

src/prime_rl/trainer/rl/train.py

@@ -53,6 +53,7 @@
     begin_backward,
     clip_grad_norm_,
     bind_process_to_gpu_numa_node,
+    configure_cpu_optimizer_threads,
     export_benchmark_json,
     filter_rl_trainer_tensor_stats_for_wandb,
     finish_backward,
@@ -118,8 +119,10 @@ def train(config: TrainerConfig):
     setup_torch_distributed(
         timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.fsdp_cpu_offload
     )
-    if config.model.optim_cpu_offload and config.model.optim_cpu_offload.numa_bind:
-        bind_process_to_gpu_numa_node()
+    if config.model.optim_cpu_offload:
+        if config.model.optim_cpu_offload.numa_bind:
+            bind_process_to_gpu_numa_node()
+        configure_cpu_optimizer_threads()
     # Configurable to support ROCm/AMD GPUs where reduced precision
     # matmul corrupts softmax over large vocabularies. Override via config
     # (e.g. matmul_precision = "highest") on ROCm.

src/prime_rl/trainer/sft/train.py

@@ -41,6 +41,7 @@
     MemoryProfiler,
     begin_backward,
     bind_process_to_gpu_numa_node,
+    configure_cpu_optimizer_threads,
     clip_grad_norm_,
     export_benchmark_json,
     finish_backward,
@@ -91,8 +92,10 @@ def train(config: SFTConfig):
     setup_torch_distributed(
         timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.fsdp_cpu_offload
     )
-    if config.model.optim_cpu_offload and config.model.optim_cpu_offload.numa_bind:
-        bind_process_to_gpu_numa_node()
+    if config.model.optim_cpu_offload:
+        if config.model.optim_cpu_offload.numa_bind:
+            bind_process_to_gpu_numa_node()
+        configure_cpu_optimizer_threads()
     # Configurable to support ROCm/AMD GPUs where reduced precision
     # matmul corrupts softmax over large vocabularies. Override via config
     # (e.g. matmul_precision = "highest") on ROCm.

src/prime_rl/trainer/utils.py

@@ -1,6 +1,7 @@
 import gc
 import heapq
 import json
+import os
 import pickle
 import shutil
 import time
@@ -423,8 +424,6 @@ def bind_process_to_gpu_numa_node() -> None:
     GPU hangs off. Must run before CPU optimizer state allocation and before the
     OMP thread pool spins up.
     """
-    import os
-
     import pynvml
 
     logger = get_logger()
@@ -453,6 +452,17 @@ def bind_process_to_gpu_numa_node() -> None:
     logger.info(f"Bound rank with GPU {device_id} to NUMA node {numa_node} ({len(cpus)} CPUs)")
 
 
+def configure_cpu_optimizer_threads() -> None:
+    available = os.sched_getaffinity(0)
+    fair_share = (os.cpu_count() or len(available)) // get_world().local_world_size
+    threads = max(1, min(len(available), fair_share))
+    torch.set_num_threads(threads)
+    get_logger().info(
+        f"CPU optimizer uses {threads} intra-op threads "
+        f"({len(available)} CPUs in this rank's affinity mask, {get_world().local_world_size} local ranks)"
+    )
+
+
 def setup_torch_distributed(timeout: timedelta = DEFAULT_TIMEOUT, enable_gloo: bool = False):
     device_id = get_world().local_rank
     torch.cuda.set_device(device_id)