[perf, trainer, training_utils, ray, worker] fix: Add set_numa_affinity() for engine workers: TrainingWorker. by sheilaliuxl · Pull Request #5627 · verl-project/verl

sheilaliuxl · 2026-03-16T21:47:56Z

What does this PR do?

Add set_numa_affinity() for engine workers: TrainingWorker.

We find a consistent performance gap of sft_trainer.py vs sft_trainer_ray.py on SFT jobs on AWS p5 instances, and this PR reduces the gap though still not fully resolved

Before PR:

Gap is 0.2292 ==> 0.2064 (-9.95%)

+ python /fsx/ubuntu/users/sliuxl/tb.py --files '../*run-ray-08-????-j*/events*' --metric train/mfu --format 10.4f
2026-03-16 21:45:48,423 [tb.py:206] WARNING - Unknonw metric `train/mfu`!
2026-03-16 21:45:48,467 [tb.py:235] INFO - Set benchmark with index 0
2026-03-16 21:45:48,468 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101684-run-ray-08-base-j-01-seq04k-ns01-sp01-bs01-pad1-spmd--20260316.125636/events.out.tfevents.1773691117.ip-10-4-145-27.4009372.0: (avg (std), med) = (    0.2292 (    0.0052),     0.2301)  # train/mfu
2026-03-16 21:45:48,512 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101685-run-ray-08-base-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.125641/events.out.tfevents.1773691587.ip-10-4-145-27.4013100.0: (avg (std), med) = (    0.2064 (    0.0229),     0.2122)  # train/mfu

Checklist Before Starting

Search for similar PRs. Paste at least one query link here: https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+numa
- [perf, megatron] chore: bind NUMA #3471 adds the util function for Megatron
Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
- {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward, fully_async, one_step_off
- If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
- {type} is in feat, fix, refactor, chore, test
- If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
- Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

After PR

1. `train/mfu`

Metric train/mfu improvement, dropping first 10 steps and last 3 steps:

avg: 0.2064 ==> 0.2189 (+6.1%)
std: 0.0229 ==> 0.0081 (-64.6%) or improvement of 2.8x

+ python /fsx/ubuntu/users/sliuxl/tb.py --files '../*run-ray-0[89]*ray*/events*' --metric train/mfu --format 10.4f
2026-03-16 21:38:17,874 [tb.py:206] WARNING - Unknonw metric `train/mfu`!
2026-03-16 21:38:17,921 [tb.py:235] INFO - Set benchmark with index 0
2026-03-16 21:38:17,921 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101685-run-ray-08-base-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.125641/events.out.tfevents.1773691587.ip-10-4-145-27.4013100.0: (avg (std), med) = (    0.2064 (    0.0229),     0.2122)  # train/mfu
2026-03-16 21:38:17,965 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101705-run-ray-09-numa-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.141110/events.out.tfevents.1773696080.ip-10-4-145-27.18170.0: (avg (std), med) = (    0.2189 (    0.0081),     0.2204)  # train/mfu

2. All Metrics

+ python /fsx/ubuntu/users/sliuxl/tb.py --files '../*run-ray-0[89]*ray*/events*' --metric '*' --format 10.4f
2026-03-16 21:38:26,932 [tb.py:206] WARNING - Unknonw metric `*`!
2026-03-16 21:38:26,932 [tb.py:216] INFO - 

[00/02] File ../sft_ray_config_pbtxt-ce3b5040--101685-run-ray-08-base-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.125641/events.out.tfevents.1773691587.ip-10-4-145-27.4013100.0:
2026-03-16 21:38:26,978 [tb.py:235] INFO - Set benchmark with index 0
2026-03-16 21:38:26,979 [tb.py:264] INFO - [243/256] : (avg (std), med) = (32768.0000 [     0.0000] (    0.0000), 32768.0000)  # train/global_tokens
2026-03-16 21:38:26,979 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.2974 [     0.0000] (    0.1332),     0.2682)  # train/grad_norm
2026-03-16 21:38:26,979 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.0199 [     0.0000] (    0.0047),     0.0200)  # train/loss
2026-03-16 21:38:26,980 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.0000 [     0.0000] (    0.0000),     0.0000)  # train/lr
2026-03-16 21:38:26,980 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.2064 [     0.0000] (    0.0229),     0.2122)  # train/mfu
2026-03-16 21:38:26,980 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.0043 [     0.0000] (    0.0023),     0.0043)  # train/total_tokens(B)
2026-03-16 21:38:26,980 [tb.py:264] INFO - [00/01] : (avg (std), med) = (    0.0000 [     0.0000] (    0.0000),     0.0000)  # val/loss
2026-03-16 21:38:26,980 [tb.py:216] INFO - 

[01/02] File ../sft_ray_config_pbtxt-ce3b5040--101705-run-ray-09-numa-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.141110/events.out.tfevents.1773696080.ip-10-4-145-27.18170.0:
2026-03-16 21:38:27,023 [tb.py:264] INFO - [243/256] : (avg (std), med) = (32768.0000 [     0.0000] (    0.0000), 32768.0000)  # train/global_tokens
2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.8776 [+    0.5802] (    9.1371),     0.2680)  # train/grad_norm
2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.0199 [+    0.0000] (    0.0046),     0.0201)  # train/loss
2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.0000 [     0.0000] (    0.0000),     0.0000)  # train/lr
2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.2189 [+    0.0125] (    0.0081),     0.2204)  # train/mfu
2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = (    0.0043 [     0.0000] (    0.0023),     0.0043)  # train/total_tokens(B)
2026-03-16 21:38:27,024 [tb.py:264] INFO - [00/01] : (avg (std), med) = (    0.0000 [     0.0000] (    0.0000),     0.0000)  # val/loss

API and Usage Example

N.A.

Design & Code Changes

Similar to #3471

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

Read the Contribute Guide.
Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
Add / Update the documentation.
Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

gemini-code-assist

Code Review

This pull request introduces a call to set_numa_affinity() in the TrainingWorker to improve performance by setting NUMA affinity. While this is a positive change for GPU environments, the review identifies a critical limitation: the underlying implementation is a no-op on NPU systems, meaning the performance fix does not apply to Ascend hardware. Furthermore, the existing code and TODO comments suggest a misunderstanding of what's required for NPU support, as fixing it would require more than just removing an early return. A comment has been added to highlight this high-severity issue.

sheilaliuxl · 2026-03-16T21:53:49Z

Adding @eric-haibin-lin @vermouth1992 @ZihengJiang for review, thanks.

sheilaliuxl · 2026-03-16T22:13:14Z

Adding @wuxibin89 as well, as they reviewed #3471.

ETOgaosion · 2026-03-18T03:18:14Z

@sheilaliuxl Thanks for implemenation! Seems that previous training worker function not migrated to engine worker.

Wondering what backend you tested? Do FSDP and Megatron both have such effect?

sheilaliuxl · 2026-03-18T06:46:39Z

Thanks for the review!

I looked at [perf, megatron] chore: bind NUMA #3471, and I believe it addresses on the Megatron side only
Therefore on my side, I've been testing with FSDP as you can guess, as shown in my full cmd below:

ubuntu@ip-10-4-201-161:$ realpath events*  # NUMA run
/fsx/....../sft_ray_config_pbtxt-ce3b5040--101705-run-ray-09-numa-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.141110/events.out.tfevents.1773696080.ip-10-4-145-27.18170.0


ubuntu@ip-10-4-201-161:$ tail -50 run_verl.sh 

TIMESTAMP=$(TZ=America/Los_Angeles date +%Y%m%d.%H%M%S)
HYDRA_FULL_ERROR=1 python \
    -m verl.trainer.sft_trainer_ray \
    checkpoint.save_contents=[] \
    data.balance_dp_token=False \
    data.max_length=4096 \
    data.micro_batch_size_per_gpu=1 \
    data.train_batch_size=8 \
    data.train_files=/fsx/....../train_multiturn_2048.parquet \
    data.train_max_samples=-1 \
    data.truncation=right \
    data.use_shm=False \
    data.val_files=/fsx/....../test_multiturn_2048.parquet \
    data.val_max_samples=-1 \
    engine.model_dtype=fp32 \
    engine.param_offload=False \
    engine.strategy=fsdp2 \
    engine.ulysses_sequence_parallel_size=1 \
    engine.wrap_policy.min_num_params=0 \
    model.enable_gradient_checkpointing=True \
    model.lora_alpha=16 \
    model.lora_rank=0 \
    model.path=/fsx/....../Qwen2.5-3B-Instruct \
    model.target_modules=all-linear \
    model.trust_remote_code=False \
    model.use_liger=False \
    model.use_remove_padding=False \
    model.use_shm=False \
    optim.betas="[0.9, 0.95]" \
    optim.clip_grad=1.0 \
    optim.lr=1e-5 \
    optim.lr_scheduler_type=cosine \
    optim.lr_warmup_steps_ratio=0.1 \
    optim.weight_decay=0.01 \
    trainer.default_local_dir=/fsx/....../sft_ray_config_pbtxt-ce3b5040--${SLURM_JOB_ID}-run-ray-09-numa-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.141110/ \
    trainer.device=cuda \
    trainer.experiment_name=sft_ray_config_pbtxt-ce3b5040--${SLURM_JOB_ID}-run-ray-09-numa-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.141110 \
    trainer.logger="[console, mlflow, tensorboard]" \
    trainer.n_gpus_per_node=8 \
    trainer.nnodes=1 \
    trainer.project_name=sliuxl-mverl--qwen2d5_3b--gsm8k_2048 \
    trainer.resume_mode=auto \
    trainer.save_freq=-1 \
    trainer.seed=1 \
    trainer.test_freq=-1 \
    trainer.total_epochs=1 \
    2>&1 | tee /fsx/....../train_log.txt

TIMESTAMP_FINISH=$(TZ=America/Los_Angeles date +%Y%m%d.%H%M%S)

…ty() for engine workers: TrainingWorker. (verl-project#5627) ### What does this PR do? Add `set_numa_affinity()` for engine workers: `TrainingWorker`. - We find a consistent performance gap of `sft_trainer.py` vs `sft_trainer_ray.py` on SFT jobs on AWS `p5` instances, and this PR **reduces** the gap though still **not** fully resolved **Before PR**: - Gap is `0.2292 ==> 0.2064 (-9.95%)` ``` + python /fsx/ubuntu/users/sliuxl/tb.py --files '../*run-ray-08-????-j*/events*' --metric train/mfu --format 10.4f 2026-03-16 21:45:48,423 [tb.py:206] WARNING - Unknonw metric `train/mfu`! 2026-03-16 21:45:48,467 [tb.py:235] INFO - Set benchmark with index 0 2026-03-16 21:45:48,468 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101684-run-ray-08-base-j-01-seq04k-ns01-sp01-bs01-pad1-spmd--20260316.125636/events.out.tfevents.1773691117.ip-10-4-145-27.4009372.0: (avg (std), med) = ( 0.2292 ( 0.0052), 0.2301) # train/mfu 2026-03-16 21:45:48,512 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101685-run-ray-08-base-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.125641/events.out.tfevents.1773691587.ip-10-4-145-27.4013100.0: (avg (std), med) = ( 0.2064 ( 0.0229), 0.2122) # train/mfu ``` ### Checklist Before Starting - [x] Search for similar PRs. Paste at least one query link here: https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+numa * verl-project#3471 adds the util function for `Megatron` - [x] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) - `{modules}` include `fsdp`, `megatron`, `veomni`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`, `fully_async`, `one_step_off` - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` ### Test **After PR** #### 1. `train/mfu` Metric `train/mfu` improvement, dropping first `10` steps and last `3` steps: 1. `avg`: `0.2064 ==> 0.2189 (+6.1%)` 2. `std`: `0.0229 ==> 0.0081 (-64.6%)` or improvement of `2.8x` ``` + python /fsx/ubuntu/users/sliuxl/tb.py --files '../*run-ray-0[89]*ray*/events*' --metric train/mfu --format 10.4f 2026-03-16 21:38:17,874 [tb.py:206] WARNING - Unknonw metric `train/mfu`! 2026-03-16 21:38:17,921 [tb.py:235] INFO - Set benchmark with index 0 2026-03-16 21:38:17,921 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101685-run-ray-08-base-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.125641/events.out.tfevents.1773691587.ip-10-4-145-27.4013100.0: (avg (std), med) = ( 0.2064 ( 0.0229), 0.2122) # train/mfu 2026-03-16 21:38:17,965 [tb.py:264] INFO - [243/256] ../sft_ray_config_pbtxt-ce3b5040--101705-run-ray-09-numa-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.141110/events.out.tfevents.1773696080.ip-10-4-145-27.18170.0: (avg (std), med) = ( 0.2189 ( 0.0081), 0.2204) # train/mfu ``` #### 2. All Metrics ``` + python /fsx/ubuntu/users/sliuxl/tb.py --files '../*run-ray-0[89]*ray*/events*' --metric '*' --format 10.4f 2026-03-16 21:38:26,932 [tb.py:206] WARNING - Unknonw metric `*`! 2026-03-16 21:38:26,932 [tb.py:216] INFO - [00/02] File ../sft_ray_config_pbtxt-ce3b5040--101685-run-ray-08-base-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.125641/events.out.tfevents.1773691587.ip-10-4-145-27.4013100.0: 2026-03-16 21:38:26,978 [tb.py:235] INFO - Set benchmark with index 0 2026-03-16 21:38:26,979 [tb.py:264] INFO - [243/256] : (avg (std), med) = (32768.0000 [ 0.0000] ( 0.0000), 32768.0000) # train/global_tokens 2026-03-16 21:38:26,979 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.2974 [ 0.0000] ( 0.1332), 0.2682) # train/grad_norm 2026-03-16 21:38:26,979 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.0199 [ 0.0000] ( 0.0047), 0.0200) # train/loss 2026-03-16 21:38:26,980 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.0000 [ 0.0000] ( 0.0000), 0.0000) # train/lr 2026-03-16 21:38:26,980 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.2064 [ 0.0000] ( 0.0229), 0.2122) # train/mfu 2026-03-16 21:38:26,980 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.0043 [ 0.0000] ( 0.0023), 0.0043) # train/total_tokens(B) 2026-03-16 21:38:26,980 [tb.py:264] INFO - [00/01] : (avg (std), med) = ( 0.0000 [ 0.0000] ( 0.0000), 0.0000) # val/loss 2026-03-16 21:38:26,980 [tb.py:216] INFO - [01/02] File ../sft_ray_config_pbtxt-ce3b5040--101705-run-ray-09-numa-j-02-seq04k-ns01-sp01-bs01-pad1-ray--20260316.141110/events.out.tfevents.1773696080.ip-10-4-145-27.18170.0: 2026-03-16 21:38:27,023 [tb.py:264] INFO - [243/256] : (avg (std), med) = (32768.0000 [ 0.0000] ( 0.0000), 32768.0000) # train/global_tokens 2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.8776 [+ 0.5802] ( 9.1371), 0.2680) # train/grad_norm 2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.0199 [+ 0.0000] ( 0.0046), 0.0201) # train/loss 2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.0000 [ 0.0000] ( 0.0000), 0.0000) # train/lr 2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.2189 [+ 0.0125] ( 0.0081), 0.2204) # train/mfu 2026-03-16 21:38:27,024 [tb.py:264] INFO - [243/256] : (avg (std), med) = ( 0.0043 [ 0.0000] ( 0.0023), 0.0043) # train/total_tokens(B) 2026-03-16 21:38:27,024 [tb.py:264] INFO - [00/01] : (avg (std), med) = ( 0.0000 [ 0.0000] ( 0.0000), 0.0000) # val/loss ``` ### API and Usage Example N.A. ### Design & Code Changes Similar to verl-project#3471 ### Checklist Before Submitting > [!IMPORTANT] > Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. - [x] Read the [Contribute Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md). - [x] Apply [pre-commit checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` - [x] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs). - [x] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... - [x] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ). (If not accessible, please try [the Feishu group (飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).) - [x] If your PR is related to the `recipe` submodule, please also update the reference to the submodule commit via `git submodule update --remote` or `cd recipe && git pull origin main`.

Add set_numa_affinity() for engine workers.

1673523

gemini-code-assist Bot reviewed Mar 16, 2026

View reviewed changes

Comment thread verl/workers/engine_workers.py

sheilaliuxl mentioned this pull request Mar 17, 2026

[perf, training_utils, fsdp, ray, worker] fix: Add set_numa_affinity() for the whole RL pipeline #5637

Closed

8 tasks done

wuxibin89 approved these changes Mar 23, 2026

View reviewed changes

wuxibin89 merged commit 95457b2 into verl-project:main Mar 23, 2026
98 of 115 checks passed

Navigation Menu

Repository navigation

[perf, trainer, training_utils, ray, worker] fix: Add set_numa_affinity() for engine workers: TrainingWorker. - #5627

[perf, trainer, training_utils, ray, worker] fix: Add set_numa_affinity() for engine workers: TrainingWorker.#5627
wuxibin89 merged 1 commit into
verl-project:mainfrom
sheilaliuxl:dev-ray-numa

sheilaliuxl commented Mar 16, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

sheilaliuxl commented Mar 16, 2026

Uh oh!

sheilaliuxl commented Mar 16, 2026

Uh oh!

ETOgaosion commented Mar 18, 2026

Uh oh!

sheilaliuxl commented Mar 18, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

sheilaliuxl commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Checklist Before Starting

Test

1. train/mfu

2. All Metrics

API and Usage Example

Design & Code Changes

Checklist Before Submitting

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

sheilaliuxl commented Mar 16, 2026

Uh oh!

sheilaliuxl commented Mar 16, 2026

Uh oh!

ETOgaosion commented Mar 18, 2026

Uh oh!

sheilaliuxl commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

sheilaliuxl commented Mar 16, 2026 •

edited

Loading

1. `train/mfu`

sheilaliuxl commented Mar 18, 2026 •

edited

Loading