diff --git a/skills/training/start-run/SKILL.md b/skills/training/start-run/SKILL.md index 9c73dab920..b55d3cfe11 100644 --- a/skills/training/start-run/SKILL.md +++ b/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 diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 066b3ba0d5..efc5f6b90b 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/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. diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index e606c88ea5..9d387ad402 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/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. diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 3e898f5f3b..ef4c385075 100644 --- a/src/prime_rl/trainer/utils.py +++ b/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)