Repository navigation
[TRTLLM-16725][feat] Add selective Model Streamer loading for Kimi K3 - #20782
yijingl-nvidia wants to merge 5 commits into
Conversation
sunnyqgg
left a comment
sunnyqgg
left a comment
There was a problem hiding this comment.
This review covers only this PR's own commit, ddfeb9a ("Add selective Model Streamer loading for Kimi K3"), against its parent 255c1a0. The stacked PRs (#20711, #20717, #20781) are reviewed in their own PRs. Note that this branch is based on an older version of #20781: its current head is 3b6ac69. The Model Streamer code this commit uses is identical in both, but #20781's new head no longer contains the #20717 move that this commit builds on. Findings: 1 blocker, 2 possibly blocking, 7 required. Six of the required items are inline; the possibly-blocking items and the remaining required item are below.
- [Possibly blocking] +1 to cascade812's thread on
tensorrt_llm/_torch/models/checkpoints/model_streamer/model_loading.py:288. This commit adds one of those per-rank raises:plan_model_streamer_load()checks only this rank's experts:expected_keyscomes frommoe.local_expert_ids(modeling_kimi_k3_loader.py:626-629).- With a checkpoint missing one routed-expert tensor, only the owning rank raises the
KeyErroratmodeling_kimi_k3_loader.py:709-712. The other ranks block in the MPI allgather. - The comment at line 702 ("Fail here, before the streaming code calls other ranks") promises something the check does not deliver.
- Fix: every rank can derive all experts' key names from the config, so checking all of them makes the error the same on every rank. Otherwise, please reword the comment.
- [Possibly blocking] +1 to chienchunhung's thread on
examples/kimi_k3/README.md:108.quick_start_kimi_k3.py:63hard-codesenable_attention_dp=True, so--tp-size 8runs DEP8, not TEP8. The README, the deployment guide (TEP8 = attention DP off) and the description's experiment all use TEP8. The repo's own evidence on whether DEP8 fits on GB300 conflicts:- The README's own numbers (114 GB replicated + 181 GB of experts per rank,
README.md:263-267) exceed 288 GB. examples/kimi_k3/disagg/README.md:32-37calls a GB300 fit "likely negative".test_kimi3.py::test_gsm8k_nvfp4_dep8runs DEP8 only with FP8 MoE-MLP reads, because the default BF16 MLP "leave[s] DEP8 no KV-cache room".- Fix: add an attention-DP switch, as
run_eval_kimi_k3.sbatchdoes with--parallel dep|tep, or drop the TEP8 sentence.
- The README's own numbers (114 GB replicated + 181 GB of experts per rank,
- [Required] The PR description does not match the code:
- It says "0 tests changed", but the commit also changes the
_model_streamer_loadhelper that existing tests intest_model_loader_checkpoint_io.pyuse, and refactorsload_mla_checkpointintest_kimi_k3_checkpoint.py. - It describes language-only loads for Kimi K3, but the commit also makes Kimi K2.5 honor
disable_mm_encoder(modeling_kimi_k25.py:1632, 1748, 1791). With the flag, K2.5 now skips the vision tower and rejects image requests.trtllm-bench throughputsets the flag for every text-only run (bench/benchmark/throughput.py:436-446), so K2.5/K3 benchmark numbers will shift because the KV-cache pool grows. Please describe this (or move the K2.5 part to its own PR) and add a K2.5 test for it. - The CodeRabbit block describes the whole 99-file stack and lists
tensorrt_llm/_torch/pyexecutor/weight_load_strategies/*files that do not exist at this head. It would end up in the squash commit message.
- It says "0 tests changed", but the commit also changes the
| mm_token_ids = None | ||
|
|
||
| if len(multimodal_params) > 0: | ||
| if self.model_config.disable_mm_encoder: |
There was a problem hiding this comment.
[Blocker] With disable_mm_encoder: true (the README's Model Streamer YAML sets it), image and video requests are rejected only here, inside the batched forward(). Nothing rejects them earlier:
- the Kimi input processor still preprocesses the media;
- the LLM API and
trtllm-servedo not check the flag (only gRPC hides media,grpc/openengine/capabilities.py:80).
A forward-time exception fails the whole batch:
PyExecutor._forward_step(py_executor.py:8973-8991) calls_handle_errors(error_msg)withoutrequests. That fails every active request and charges the error budget (error_classification.py:118-148), so a few image requests in quick succession shut the executor down.- Under attention DP, which
quick_start_kimi_k3.py:63hard-codes, only the rank that received the image raises, beforeself.llm.forward. The other ranks enter the MoE collectives and wait for it.
The executor already has a per-request hook for this:
_validate_requestcallsmodel.validate_multimodal_request_data(mm_data)at admission (py_executor.py:6876-6884).- A failure there rejects only that request, with
charge_budget=False(py_executor.py:7340-7343). - The hook's docstring (
modeling_multimodal_mixin.py:470-478) says that conditions a model would reject insideforwardbelong there.
Suggested fix:
- Implement
validate_multimodal_request_data(self, mm_data)onKimiK25ForConditionalGeneration; K3 inherits it. It should raiseValueErrornamingdisable_mm_encoder=Truewhen the flag is set and the request carries multimodal data, the same condition the forward check uses. - Keep the forward check as a backstop.
- Add a test that sends one image request next to a text request and checks that only the image request fails.
| every serving rank. Choose the `aarch64` or `x86_64` wheel for the worker CPU: | ||
|
|
||
| ```bash | ||
| python -m pip install /path/to/runai_model_streamer-*.whl |
There was a problem hiding this comment.
[Required] This install step contradicts the rest of the README and has no version:
- The Slurm jobs see only packages in the mounted
.venv-3.12. README.md:49-57 says "Packages installed only in the preparation container's base environment are not" available, andfla-coreis installed with.venv-3.12/bin/python -m pip(README.md:68). - No package version is given, but the code needs
stream_files(tensor_names=..., owned=...), which 0.16.x does not have.
Following it literally, every rank fails with the ImportError or TypeError from stream_batched_tensors. That happens only after the model is built and the ranks have gathered tensor names, so a whole multi-node allocation is wasted. runai-model-streamer 0.17.0 and 0.17.1 are on PyPI (the description's run used 0.17.1), so this can be:
.venv-3.12/bin/python -m pip install "runai-model-streamer>=0.17"|
|
||
| The streamer reads only tensors needed by the ranks on each node and uses | ||
| node-local distributed loading when multiple ranks share a node. Its current | ||
| scope is local Hugging Face safetensors checkpoints for the text Kimi K3 model. |
There was a problem hiding this comment.
[Required] "Its current scope is ... the text Kimi K3 model" contradicts this commit:
KimiK3ForConditionalGeneration.plan_model_streamer_load(modeling_kimi_k3_vl.py:550-590) also streams the vision tower and the projector.- The quick-start path this section recommends (
--checkpoint-io-policy model_streamer, nodisable_mm_encoder) takes exactly that vision path with the multimodal checkpoint. - Putting
disable_mm_encoder: truein the Model Streamer YAML makes it look required.
Please state that both the text-only and the multimodal Kimi K3 classes are supported, or reject the VL plan if the vision path has not been validated end to end. Show disable_mm_encoder as a separate, optional setting.
| # ------------- Kimi K3 (KimiLinear) unit tests --------------- | ||
| - unittest/_torch/modeling/test_kimi_k3_checkpoint.py TIMEOUT (15) | ||
| - unittest/_torch/modeling/test_kimi_k3_mla.py TIMEOUT (15) | ||
| - unittest/_torch/models/test_kimi_k3_model_streamer.py |
There was a problem hiding this comment.
[Required] All four cases in test_kimi_k3_model_streamer.py run without a GPU:
- the multimodal tests use fake parts;
- the expert test's harness returns a CPU
torch.empty(0)fromparameters(); run_model_streamer(..., world_size=1)skips the distributed setup, andstream_batched_tensorsis monkeypatched.
tests/AGENTS.md says "CPU-only tests go in l0_cpu.yml", and the CPU stages already import modeling_kimi_linear for other cpu_only Kimi tests. BowenFu asked for this entry to get the expert path into CI; l0_cpu.yml does that too, without using a B200 slot, and also runs on arm.
Please add pytestmark = pytest.mark.cpu_only to the file and move the entry to l0_cpu.yml. The CPU stage runs -m cpu_only, and tests/unittest/conftest.py skips files without the marker.
| "that never send image/video/audio inputs; such requests are " | ||
| "rejected. Only takes effect for model implementations that support " | ||
| "it (currently Mistral3 and the Qwen3-VL / Qwen3.5-VL family); a " | ||
| "it (currently Kimi K2.5/K3, Mistral3, and the Qwen3-VL / " |
There was a problem hiding this comment.
[Required] This sentence is being rewritten, and it is still incomplete. DeepSeek-V4.1 multimodal (modeling_deepseekv41_multimodal.py:60) and the GLM-5 Next vision model (modeling_glm5_next_vision.py:1054) also honor disable_mm_encoder. The field description, the regenerated serve schemas and the comment in model_config.py therefore tell their users the flag is a no-op.
Please complete the list, or replace it with "models whose implementation checks ModelConfig.disable_mm_encoder". Then regenerate the schemas with scripts/generate_trtllm_serve_schemas.py.
| return model | ||
|
|
||
|
|
||
| def test_multimodal_wrapper_adds_vision_group(monkeypatch): |
There was a problem hiding this comment.
[Required] CODING_GUIDELINES.md says: "Always annotate functions. Make the return type None if the function does not return anything".
- None of the functions, helpers or fake methods in this new file are annotated, nor is the nested
streamintest_kimi_k3_checkpoint.py:285. test_model_streamer_matches_regular_loaderin the same commit is annotated.- It matters here because the fakes stand in for
stream_batched_tensorsand for the expert-loader signature, which already drifted once in this PR (get_optional_tensor).
Please annotate them as test_model_loader_checkpoint_io.py does: monkeypatch: pytest.MonkeyPatch, -> None, consume: Callable[[str, torch.Tensor], None].
| load_function)``. | ||
|
|
||
| Raises: | ||
| - If the checkpoint lacks a tensor that this rank needs (KeyError). |
There was a problem hiding this comment.
[Required] CODING_GUIDELINES.md asks for Google-style docstrings. Please write this one, and the one at modeling_kimi_k3_vl.py:568-570, as KeyError: If the checkpoint lacks a tensor that this rank needs., the way model_streamer/model_loading.py:132-135 does.
Optional style points in the same code:
_ROUTED_EXPERT_KEYis a compiled regex;_ROUTED_EXPERT_KEY_REwould say so.groupscould carry a type annotation instead of the# groups:comment.- The new signatures mix
Set/Dictwithset/dict.
c9bdeb0 to
db0aa91
Compare
sunnyqgg
left a comment
sunnyqgg
left a comment
There was a problem hiding this comment.
Follow-up to my review of 2026-11-03: 16 of the 19 lower-priority points I left out of it still apply at db0aa91. The other 3 are fixed or no longer apply. They are marked [Nice to have] and none blocks the merge on its own.
- [Nice to have] There is no
/bot runon this PR yet, and the description lists the new tests without results, so nothing shows that they run and pass. Please post/bot runonce the stack allows and link the run. It should show the 8test_model_streamer_matches_regular_loadercases and the 4test_kimi_k3_model_streamer.pycases passing rather than skipped (the expert test callspytest.importorskip("fla")), and the CPU stage passing the changed_model_streamer_loadhelper intest_model_loader_checkpoint_io.py.
|
|
||
| def load_weights(self, weights) -> None: | ||
| """Load vision + projector + LLM weights from checkpoint.""" | ||
| self._ensure_mm_encoder() |
There was a problem hiding this comment.
[Nice to have] load_weights() now calls self._ensure_mm_encoder() first, but test_language_weights_preserve_checkpoint_dir (test_modeling_kimi_k25.py:269-273) calls it on a SimpleNamespace without that method, so the test now fails with AttributeError. CI stays green only because that file has no cpu_only marker and is in no test list. It is the only test of the checkpoint_dir/checkpoint_prefix pass-through that K3's per-file expert loading needs (modeling_kimi_k3_loader.py:1247-1251). Please add _ensure_mm_encoder=lambda: None to the stand-in and mark the test pytest.mark.cpu_only, so the unittest/_torch/modeling entry in l0_cpu.yml:68 runs it.
| self.loaded_weights = weights | ||
|
|
||
| monkeypatch.setattr(modeling_kimi_k25, "_has_meta_tensors", lambda model: False) | ||
| model = _multimodal_wrapper(TextModel(), VisionModel()) |
There was a problem hiding this comment.
[Nice to have] Both VL plan tests make _ensure_mm_encoder() a no-op: this one passes a ready VisionModel(), and the text-only test sets disable_mm_encoder=True (test_kimi_k3_model_streamer.py:83). In production, MetaInitMode leaves mm_encoder as None until the plan rebuilds it (modeling_kimi_k3_vl.py:574), so removing that call keeps these tests green while multimodal streamer loads lose the vision group. The KeyError at modeling_kimi_k3_vl.py:588 is untested too. Please add a case that starts from mm_encoder=None with a fake _VISION_MODEL_CLS and _vlm_pretrained_config and checks that the last group loads into the rebuilt encoder, plus a no-vision-keys case that expects the KeyError.
| ) | ||
| assert sorted(regular_events) == expected_events | ||
| assert sorted(streamed_events) == expected_events | ||
| assert [len({_expert_layer(name) for name in batch}) for batch in streamed_batches] == [ |
There was a problem hiding this comment.
[Nice to have] This test never checks which tensors were streamed: it compares what the fake loader read, the prepare/finalize events and the layers per batch. The fake loader (test_kimi_k3_model_streamer.py:133-143) reads only its own expert, so a plan whose groups also named expert 1 (owned by another rank) would still pass, although each node would then read experts its ranks do not own, which defeats the selective load. Please assert the exact set, e.g. set().union(*streamed_batches) == expected, where expected holds experts 2 and 0 x w1/w2/w3 x kinds for every layer, plus expert 2's w1.pre_quant_scale.
| def build() -> KimiLinearForCausalLM: | ||
| with torch.device("cuda"): | ||
| model = KimiLinearForCausalLM( | ||
| _config(fp8, mapping=mapping, prefix=prefix, deferred=False, cute_bmm=True) |
There was a problem hiding this comment.
[Nice to have] This test builds only deferred=False, cute_bmm=True, and FP8 only with weight_scale_inv (test_kimi_k3_checkpoint.py:255). So MLA never creates k_b_proj_trans_dequant/v_b_proj_dequant (mla.py:802), although that is the default SM100 path (use_cute_dsl_blockscaling_bmm defaults to False, llm_args.py:7232); the loop at test_kimi_k3_checkpoint.py:310-315 compares None with None, and the weight_scale branch of tensor_and_scale_names (modeling_kimi_k3_loader.py:721) never runs. Please reuse the sibling test's (deferred, cute_bmm) pairs and three (fp8, scale_name) pairs (test_kimi_k3_checkpoint.py:126-129); the trailing loop can then go, since named_parameters() covers those parameters.
| # not be kept for runtime forward. | ||
| self.mm_encoder = None | ||
| if not DISAGG: | ||
| if not DISAGG and not model_config.disable_mm_encoder: |
There was a problem hiding this comment.
[Nice to have] No test builds the real KimiK25ForConditionalGeneration or KimiK3ForConditionalGeneration with disable_mm_encoder=True; the streamer tests use a wrapper that skips __init__ and forward. If this guard (or modeling_kimi_k3_vl.py:498) is reverted, a model built without MetaInitMode still builds and loads the vision tower, and without the check at modeling_kimi_k25.py:1791 an image request fails on None.forward, with no test noticing. Mistral3 has such a test (test_modeling_mistral.py:182). Please add construction tests for both classes with the flag, with and without MetaInitMode (no mm_encoder.* parameters, encoder still None after load_weights), and a forward test that expects the rejection.
| self._load_trunk_params, | ||
| params={name: params[name]}, | ||
| name_map={name: key}, | ||
| log_summary=False, |
There was a problem hiding this comment.
[Nice to have] The streamed load drops diagnostics that the native load prints: the plan never runs _validate_checkpoint_keys, so the unmatched-key warning (modeling_kimi_k3_loader.py:916-920) is gone, and with log_summary=False on every group the KV-B summary (:1124-1127) is never logged. A checkpoint with stray tensors therefore warns under native and is silent under model_streamer. K2.5/K3 also log nothing when disable_mm_encoder skips the encoder (modeling_kimi_k25.py:1632, modeling_kimi_k3_vl.py:498), unlike modeling_qwen3vl.py:1565-1569. Please run the unmatched-key check from the plan, log the KV-B summary once in finish(), and add the same info line when the encoder is skipped.
| load runs `run_model_streamer()` for real. The returned list records the names | ||
| that the fake streamer delivered. | ||
| """ | ||
| from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM |
There was a problem hiding this comment.
[Nice to have] This generic ModelLoader helper now imports the Kimi K3 model, and through kimi_kda_mixer.py:27 the fla package, only to spec a MagicMock, so any Kimi/KDA/fla import problem breaks unrelated loader tests. A small local nn.Module subclass that defines load_weights and plan_model_streamer_load would serve as the spec. Also, test_kimi_k3_model_streamer.py calls pytest.importorskip("fla") only in the expert test (test_kimi_k3_model_streamer.py:101), but its two multimodal tests import the same chain and would error instead of skip; please move it to module level, and consider putting that file in tests/unittest/_torch/modeling/ with the other Kimi tests.
| self._ensure_mm_encoder() | ||
| if self.mm_encoder is not None: | ||
| # Name prefixes of vision encoder and projector tensors. | ||
| vision_prefixes = ( |
There was a problem hiding this comment.
[Nice to have] These prefixes copy the _detect_prefix candidates in KimiK25VisionModel.load_weights (modeling_kimi_k25.py:910 and :917). If the vision loader learns a new prefix and this tuple is not updated, the streamed load skips those tensors or raises the KeyError at modeling_kimi_k3_vl.py:588, while the native load works. Please hoist the two candidate lists into module constants in modeling_kimi_k25.py and build this tuple from them.
| name_map: Dict[str, str], | ||
| log_summary: bool = True, | ||
| ) -> int: | ||
| """Load every non-expert trunk parameter concurrently (with the |
There was a problem hiding this comment.
[Nice to have] Several docstrings no longer match the code. This one says it loads "every non-expert trunk parameter concurrently", but the plan passes one parameter; _load_attention_fp8 (modeling_kimi_k3_loader.py:386) does not document projection_names; the class docstring (:335-342) is vague ("mixin-like", "would read") where it should say the methods work only as a base of KimiLinearForCausalLM. In modeling_kimi_k3_vl.py, :458-461 says "only __init__ differs", :495-496 says the encoder is recreated in load_weights (the plan now does it too), and :557-558 says every rank keeps the whole vision encoder, which is false when the tower is TP-sharded (modeling_kimi_k25.py:383-385). Please update them.
| key for key in checkpoint_keys if key.startswith(vision_prefixes) | ||
| } | ||
| if not vision_keys: | ||
| raise KeyError("Kimi K3 Model Streamer found no vision weights") |
There was a problem hiding this comment.
[Nice to have] With checkpoint_io_policy=model_streamer, a checkpoint that declares a vision tower but stores no vision or projector tensors fails here on every rank, while the native load only warns and continues (modeling_kimi_k25.py:923-925). Raising is safer than serving an uninitialized tower, but the message gives the user no way forward. Please list the searched prefixes and suggest disable_mm_encoder: true, e.g. KeyError(f"Kimi K3 Model Streamer found no tensors starting with {vision_prefixes}; set disable_mm_encoder: true to load only the language model").
2a12f1e to
3a29fdc
Compare
2e158ca to
c8b0cc8
Compare
b09193f to
da61750
Compare
To reduce the complex code in model_loader.py, created new files to hold helper functions: - tensorrt_llm/_torch/models/checkpoints/gms/model_loading.py - tensorrt_llm/_torch/models/checkpoints/mx/model_loading.py - tensorrt_llm/_torch/models/checkpoints/memory_init.py - tensorrt_llm/_torch/models/checkpoints/model_hooks.py Refactor ModelLoader.load() to be more compact and simpler to understand. Create private helper functions in model_loader.py to hold load_format=auto loading path, which includes the HF loading, the rank-striped-read-ahead loading and custom checkpoint loading options. Remove an old duplicate doc docs/source/torch/features/checkpoint_loading.md. Clean up the checkpoint loading doc docs/source/features/checkpoint-loading.md: - Add missing checkpoint format. - Simplify explanation. - Remove extra details most users don't need to know. Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Move the weight-loading methods of KimiLinearForCausalLM and the helpers that only they use into the new file modeling_kimi_k3_loader.py. modeling_kimi_linear.py had over 3,300 lines, and about 1,100 of them only turn checkpoint tensors into model parameters. The code moves unchanged: KimiLinearForCausalLM inherits the methods from KimiK3WeightLoaderMixin, and tests import the moved helpers from the new file. The inference module _Fp8BlockScaleWeightReadLinear stays in modeling_kimi_linear.py; the loader imports it inside the functions that use it, because modeling_kimi_linear.py imports the loader. No functional change. Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Add checkpoint_io_policy="model_streamer", which loads a Hugging Face checkpoint with Run:ai Model Streamer. Model Streamer reads each byte from storage once per node and shares it among the ranks of the node, so large checkpoints load faster. A model describes what it needs in plan_model_streamer_load(), which returns a ModelStreamerPlan: groups of checkpoint tensor names, each with a function that loads the group. The new models/checkpoints/model_streamer/model_loading.py streams only the tensors that the ranks of a node need, with run_model_streamer(), and loads each group as soon as its tensors arrive. A separate draft checkpoint, such as a DSpark drafter, is read with the regular checkpoint reader. The next commit adds this for Kimi K3. Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Implement plan_model_streamer_load() for Kimi K3 and its multimodal wrapper, with the ModelStreamerPlan from the previous commit. Each rank lists only the tensors that it needs, for example its own experts, so each node streams only what its ranks use. Also support Kimi K3 language-only loads that skip the multimodal encoder, document Model Streamer loading in the Kimi K3 README, and add tests. Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
da61750 to
764e372
Compare
There was a problem hiding this comment.
My previous comments have been addressed; thanks. 2 follow-up comments:
-
Reject media requests before scheduling. With
disable_mm_encoder=True, image or video requests fail inside batchedforward(). The executor then fails all active requests, including valid text requests, and consumes its error budget. Please reject media throughvalidate_multimodal_request_data()before scheduling. Add a mixed batch of separate text and media requests, and verify that only the media request fails. See the forward check. -
Update the existing K2.5 test double. The new
_ensure_mm_encoder()call breakstest_language_weights_preserve_checkpoint_dir. ItsSimpleNamespacelacks that method, so the test raisesAttributeErrorbefore it checks checkpoint metadata. Please update the test double and run the test. See the new helper call.
|
Could you please rebase this branch onto the latest upstream |
Description
This PR adds the Model Streamer load for Kimi K3. Each rank lists only the checkpoint tensors that it needs, for example its own experts, so each node streams only what its ranks use.
It also supports Kimi K3 language-only loads, which skip the multimodal encoder (
disable_mm_encoder), and adds a Model Streamer section to the Kimi K3 README.The regular Kimi K3 load (
_load_expert_slices()) and the Model Streamer load put every routed expert into its MoE backend slot with the same two functions ofmodeling_kimi_k3_loader.py._load_expert_into_slot()prepares a layer's staging buffers on its first expert and finalizes the layer when its last slot fills._check_expert_slots()checks at the end that every slot is filled and every staged layer is finalized. So a change to these steps lands in both loads.This PR must merge after #20717 and #20781. Its branch contains the commits of #20711, #20713, #20717 and #20781. Review only the last commit: https://github.com/NVIDIA/TensorRT-LLM/pull/20782/files/c2c1c87a40e0ad51141b09ac60422669fc7250ed..764e372667365352989b60378f954b9ad3fedc88
The Model Streamer expert load also streams the optional per-expert tensors that the checkpoint stores, such as the AWQ pre-quant scale, and gives them to the Kimi K3 expert loaders through their
get_optional_tensorargument (#20823).test_streamer_expert_loading_matches_regular_loadercovers experts with and without the optional tensor.Test Coverage
4 tests added, 0 tests changed:
tests/unittest/_torch/modeling/test_kimi_k3_checkpoint.pytests/unittest/_torch/models/test_kimi_k3_model_streamer.pyThe new tests build their checkpoint with
mla_checkpoint(), which this PR splits out of the test helperload_mla_checkpoint(). It keeps the helper'srandom_gate_upoption from #20717.tests/unittest/_torch/models/test_kimi_k3_model_streamer.pynow runs in the B200 pre-merge list (l0_b200.yml).Load-time experiment
This experiment checks the Model Streamer change from the review:
stream_files(..., owned=False), and notorch.cuda.synchronize()after each streamed tensor. With the change, the weights and the outputs stay the same, and Model Streamer loads the weights about 6% faster.Setup:
quick_start_kimi_k3.py.1.4.0rc3.dev202610290007with the Python files of commitddfeb9abf9, and runai-model-streamer 0.17.1. This branch differs from that commit in the checks, errors, names and process-group setup of the Model Streamer load, and its Kimi K3 expert-load steps are in functions that the regular load shares. In the setup of this experiment, each node still streams the tensors that its 4 ranks need, and the 4 ranks share the reads.model_loading.pybefore the change (defaultowned=True, sync after each tensor). New code: this branch.Load time in seconds for the slowest of the 8 ranks, with the median rank in parentheses:
checkpoint_io_policyweight_population_secondstotal_model_loading_secondsmodel_streamer_secondsnativemodel_streamermodel_streamermodel_streamermodel_streamernativemodel_streamer_seconds, because the test hook that records this value started after run 2.Accuracy: all six runs gave the same weights. On each of the 8 ranks, the checksums of all 3,296 parameters and buffers (209 GB for each rank) are equal. Greedy outputs (temperature 0, 8 prompts, 128 tokens for each prompt) have the same token IDs in all six runs.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.PRs in this series
These PRs refactor checkpoint loading and add Model Streamer loading. They merge in this order, top to bottom.
checkpoint_io_policyfrom plain strings to the string enumCheckpointIOPolicy.CheckpointFormatfor the built-in checkpoint formats, make the checkpoint I/O helpers take onlyLoadFormat, and make the checkpoint I/O policy resolution easier to read.distributed/initialization.py.ModelLoader.load()shorter: move its MX, GMS, memory and model-hook helpers to new files inmodels/checkpoints/.BaseCheckpointLoader.modeling_kimi_linear.pytomodeling_kimi_k3_loader.py.checkpoint_io_policy="model_streamer", which loads a checkpoint with Run:ai Model Streamer.🤖 Generated with Claude Code