perf: optimize model loading speed#1465
Conversation
There was a problem hiding this comment.
Pull request overview
This PR targets faster loading of large fused MoE checkpoints by reducing the overhead from many tiny per-expert device copies, aiming to improve startup time and reduce warmup-related NCCL timeouts.
Changes:
- Added a shared helper (
_copy_expert_shard) plus a newstage_expert_weightAPI onFusedMoEto support writing expert shards into a CPU staging buffer. - Updated the model loader to batch expert tensor arrivals per fused MoE parameter, then flush to GPU with a single large copy.
- Added safety flushing and error surfacing behavior to ensure staged buffers are drained even when under-filled.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| atom/model_ops/moe.py | Adds staging support (stage_expert_weight) and consolidates expert-shard copy logic via _copy_expert_shard. |
| atom/model_loader/loader.py | Implements per-parameter staging buffers with a batched flush to reduce per-expert H2D copy overhead during load. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def _stage_task(param, full_param_name, shard_id, global_expert_id, loaded_weight): | ||
| pid = id(param) | ||
| with staging_lock: | ||
| if pid in fallback_pids: | ||
| _fallback( | ||
| param, full_param_name, shard_id, global_expert_id, loaded_weight | ||
| ) | ||
| return | ||
| entry = staging_map.get(pid) | ||
| if entry is None: | ||
| moe = _lookup_moe_module(full_param_name) | ||
| entry = staging_map[pid] = { | ||
| "staging": _make_staging(param), | ||
| "arrived": 0, | ||
| "expected": moe.expected_batched_arrivals(param), | ||
| "moe": moe, | ||
| "param": param, | ||
| "lock": threading.Lock(), | ||
| } | ||
|
|
| def expected_batched_arrivals(self, param: torch.nn.Parameter) -> Optional[int]: | ||
| w13_params = [ | ||
| getattr(self, n, None) | ||
| for n in ("w13_weight", "w13_weight_scale", "w13_bias") | ||
| ] | ||
| w2_params = [ | ||
| getattr(self, n, None) for n in ("w2_weight", "w2_weight_scale", "w2_bias") | ||
| ] | ||
| if any(param is p for p in w13_params if p is not None): | ||
| return self.local_num_experts * 2 | ||
| if any(param is p for p in w2_params if p is not None): | ||
| return self.local_num_experts | ||
| return None |
|
give some numbers for other models like dsr1/dsv4/glm5 ? |
|
This can reslove an issue that Kimi model load may randomly failed due to the slow ssds, and the broadcast ended timeouts. Or ROCm/aiter#4081 can fix this too by increasing the rccl process group timeout from 600s to 1200s. https://github.com/SemiAnalysisAI/InferenceX/actions/runs/28670767395/job/85033302286#step:4:419 Traceback (most recent call last):
File "/usr/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
self.run()
File "/usr/lib/python3.12/multiprocessing/process.py", line 108, in run
self._target(*self._args, **self._kwargs)
File "/app/ATOM/atom/model_engine/async_proc.py", line 130, in __init__
self.runners = [runner_class(rank, *args, **kwargs)]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/app/ATOM/atom/model_engine/model_runner.py", line 736, in __init__
self.warmup_model()
File "/app/ATOM/atom/model_engine/model_runner.py", line 1169, in warmup_model
self.forward(dummy_batch)
File "/opt/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/app/ATOM/atom/model_engine/model_runner.py", line 2193, in forward
fwd_output = self.postprocess(
^^^^^^^^^^^^^^^^^
File "/app/ATOM/atom/model_engine/model_runner.py", line 2112, in postprocess
sampled_tokens = get_tp_group().broadcast(sampled_tokens, src=0)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/app/aiter-test/aiter/dist/parallel_state.py", line 1018, in broadcast
torch.distributed.broadcast(
File "/opt/venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py", line 83, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/venv/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py", line 2907, in broadcast
work = group.broadcast([tensor], opts)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.distributed.DistBackendError: [2] is setting up NCCL communicator and retrieving ncclUniqueId from [0] via c10d key-value store by key '0', but store->get('0') got error: wait timeout after 600000ms, keys: /default_pg/0//3//cuda//0 |
Review: perf: optimize model loading speedNice speedup, and the shape of the change is sound. Batching per-expert H2D copies into one pinned staging buffer per fused param is the right lever, A few things worth addressing before merge, roughly in priority order: 1. Pinned buffer is allocated + zeroed under the global lock (undercuts the speedup)In entry = staging_map[pid] = {"staging": _make_staging(param), ...}
2.
|
|
Hi @valarLip, Thanks for the review! I hope the commits I just pushed address these comments:
A helper function has been added, and pinning / zeroing can now be executed in parallel.
A new environment variable,
I have changed the warning to a runtime assertion. The example in your comment does not seem likely to happen in practice. Disabling
The eligibility check logic has been refined, so they will no longer be marked as batchable.
I think this is fine because model loading only happens once and this step is fast. Zeroing the buffer can also help prevent potential correctness issues.
I have updated the docs and added performance measurements to this PR.
|
|
how is this happend? it only cost 30s for dsv4 pro for me |
Maybe the Tensorwave SSD is a little slow? |
|
Also Kimi have a problem that some rank loads maybe lagged. |
i see.. tw node have some issue... our solution was disable_mmap, you can check this one https://github.com/ROCm/ATOM/actions/runs/28809636117/job/85434147979#step:7:112 however your way looks works better |
Thanks, it makes a lot easier. We will give it a try by setting |
Yeah disable_mmap works well, I submit a run, and it can reduce the load time from https://github.com/SemiAnalysisAI/InferenceX/actions/runs/28881278580/job/85677479526 |
|
could you please fix the code style fail to make ci run |
|
I have updated the code and ruff should pass now. |
| raise RuntimeError( | ||
| f"Batched loader: {len(pending)} MoE param group(s) under-filled " | ||
| f"Set ATOM_LOADER_NUM_THREADS=1 to use the per-expert loader." | ||
| ) |
| | Variable | Type | Default | Description | | ||
| |----------|------|---------|-------------| | ||
| | **ATOM_DISABLE_MMAP** | bool | false | If set to `true`, disable memory-mapped file loading for model weights. Useful in containerized environments where mmap may cause issues. | | ||
| | **ATOM_LOADER_NUM_THREADS** | int | 16 | Worker threads for weight loading. `>1` (default `16`) enables the batched parallel loader (per-fused-param CPU staging flushed with a single H2D copy) with that many threads; set to `1` to fall back to the original sequential per-expert path. Raise on high-core hosts if loading is CPU-bound. | |
| # Worker threads for weight loading. >1 (default 16) enables the batched | ||
| # parallel loader (per-fused-param CPU staging flushed with one H2D copy) | ||
| # with that many threads; set to 1 to fall back to the original sequential | ||
| # per-expert path. | ||
| "ATOM_LOADER_NUM_THREADS": lambda: int(os.getenv("ATOM_LOADER_NUM_THREADS", "16")), |
|
Suggest also flipping the CI Rationale: with the batched parallel loader introduced here, weight loading no longer needs the non-mmap path, so CI can keep mmap enabled by default instead of forcing Proposed change in disable-mmap:
description: >-
- When 'true' (default) inject `-e ATOM_DISABLE_MMAP=true`. Set 'false' for
- callers whose container never set it (e.g. mmstar) to stay byte-for-byte
- equivalent.
+ When 'true', inject `-e ATOM_DISABLE_MMAP=true` into the container.
+ Defaults to 'false' (mmap enabled) now that the batched parallel loader
+ makes the non-mmap load path unnecessary; set 'true' for callers that
+ still need mmap disabled.
required: false
- default: "true"
+ default: "false"Blast radius (workflows using this action):
Note: |
Suggested-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com>
Thanks! I have your patch applied. |
That's strange.. Unfortunately I don't have access to MI355 machines right now, maybe I can test it again once I regain access. |
| # Worker threads for weight loading. >1 (default 16) enables the batched | ||
| # parallel loader (per-fused-param CPU staging flushed with one H2D copy) | ||
| # with that many threads; set to 1 to fall back to the original sequential | ||
| # per-expert path. | ||
| "ATOM_LOADER_NUM_THREADS": lambda: int(os.getenv("ATOM_LOADER_NUM_THREADS", "16")), |
| raise RuntimeError( | ||
| f"Batched loader: {len(pending)} MoE param group(s) under-filled " | ||
| f"Set ATOM_LOADER_NUM_THREADS=1 to use the per-expert loader." | ||
| ) |
|
Hi @valarLip , I investigated this a bit further, and it looks like this PR and Model loading mainly consists of two stages:
This PR targets the second stage, while On TensorWave/CI machines, the model may stored on NFS (wekafs). When the model is loaded through After the weights are in host memory, ATOM issues a large number of tiny H2D copies, which account for a significant portion of the loading time (as shown in this PR's description, which are all measured when model is on the local SSD). This PR batches those copies together, substantially reducing the H2D overhead. I have just reverted the Thanks! |
|
sounds not make sense to me, h2d never take that long |
|
Hi @valarLip , I tested on MI355X again, and the results can be reproduced stably. The following command was used to launch the ATOM server: AITER_BF16_FP8_MOE_BOUND=0 \
ATOM_MOE_GU_ITLV=1 \
ATOM_DISABLE_MMAP=0 \
ATOM_USE_TRITON_MXFP4_BMM=1 \
AITER_QUICK_REDUCE_QUANTIZATION=INT4 \
MORI_SHMEM_HEAP_SIZE=8G \
python -m atom.entrypoints.openai_server \
--model deepseek-ai/DeepSeek-V4-Pro \
--server-port 8000 \
--kv_cache_dtype fp8 \
--max-model-len 10240 \
--trust-remote-code \
--max-num-batched-tokens 16384 \
-tp 4 \
--gpu-memory-utilization 0.85 \
--no-enable_prefix_caching \
--no-enable_chunked_prefill \
--scheduler-delay-factor 10Here are the detailed logs: Model load takes ~880s (14m40s). The safetensors shard read finishes in ~8s, and the remaining ~14min is entirely the per-expert copy drain. Model load takes ~48s, all four ranks finish within 3s of each other (04:46:02–04:46:05). |
valarLip
added
ci:atom
and removed
ATOM
|
can you please help check why this never show any diff in our ci run, it's also tw nodes same as SA's |
I think this is expected. This PR won't help reducing load time on slow I/O scenerios (e.g. the model is located on NFS drive and (The logs attached above are captured on TensorWave's machine, where model are placed on the RAID nvme device /dev/md0) |
|
@gyohuangxin take a look this？ |
Makes sense. Our CI loads models from NFS, so the bottleneck is likely storage I/O. |
Looks like it has been solved by commit d13bfb0 |
Motivation
Loading large MoE checkpoints is slow on MI355X. Kimi-K2.5-MXFP4 stores each expert's
weight as a separate tensor, so the weight loader copies every expert into its slice of the
fused GPU parameter individually. Each copy moves very little data but still pays a fixed kernel-launch +
pageable-staging + GIL cost, and it's that per-copy overhead, not bandwidth, that dominates
load time. It shows up as a stall after the safetensors shard progress bar finishes, while
the loader thread pool drains the per-expert copies.
Technical Details
Checkpoint iteration and the loader thread pool are unchanged. Only the way a fused MoE
parameter's experts reach the GPU changes:
the buffer is flushed to the GPU with one large H2D copy once all experts have arrived
(
atom/model_loader/loader.py). This turns the hundreds of thousands of small copies intoone copy per fused param.
FusedMoE.stage_expert_weightwrites an expert's shard into the staging buffer instead ofthe GPU param.
This also resolves an issue when it takes too long to warm up the model, then NCCL times out and crashes.
Test Plan
Model: Kimi-K2.5-MXFP4
TP4 on MI355X (gfx950)
The metric is
load_model()wall time on the slowest TP rank.Test Results