Skip to content

[feat] Load Kimi K3 weights from object storage with Model Streamer - #21651

Draft
benamizilberNvidia wants to merge 14 commits into
NVIDIA:mainfrom
benamizilberNvidia:object-storage-loading-on-20782
Draft

[feat] Load Kimi K3 weights from object storage with Model Streamer#21651
benamizilberNvidia wants to merge 14 commits into
NVIDIA:mainfrom
benamizilberNvidia:object-storage-loading-on-20782

Conversation

@benamizilberNvidia

@benamizilberNvidia benamizilberNvidia commented Nov 4, 2026 •
edited
Loading

Copy link
Copy Markdown

@coderabbitai summary

Description

Stacked on #20782. This branch contains the 9 commits of #20782 plus 5 commits of its own. Until #20782 merges, review only the last 5 commits: compare view. Once #20782 merges, I will rebase onto main and the PR shows only those 5 commits. It is a draft until then. The conflict with main (modeling_kimi_linear.py) comes from #20782, not from these commits.

With checkpoint_io_policy: model_streamer, --model can be an s3://, gs:// or az:// URI. The weights stream from the bucket through Run:ai Model Streamer and are never copied to local disk.

The 5 commits:

  1. Add object storage helpers for model metadata: new llmapi/object_storage.py. It detects object-storage URIs and copies the files that Hugging Face loaders read from local disk (config, tokenizer, remote code, not the weights) into a node-local cache directory named after a hash of the URI. A .runai_complete marker makes later pulls no-ops, and concurrent processes on a node wait for the first one. TRTLLM_OBJECT_STORAGE_FORCE_PULL=1 downloads the files again once per launch.
  2. Accept object-storage model URIs in the LLM API: an object-storage model requires checkpoint_io_policy="model_streamer" and any other policy raises before anything is downloaded. A draft model or LoRA directory on object storage is rejected, a local one is allowed. The config, generation config and tokenizer are read from the local metadata copy, pulled once per node by local rank 0.
  3. Stream Kimi K3 weights from object storage: ModelLoader hands the URI, and the local metadata directory, to Model Streamer. The streamer's tensor source map reads model.safetensors.index.json from the local copy and maps each tensor to its file URI in the bucket.
  4. Document loading Kimi K3 from object storage: the URI forms, the required policy, the metadata cache and its force-pull option, and the limits, in the Kimi K3 README and the checkpoint loading guide.
  5. Add a runai extra for loading from object storage: a new optional extra in setup.py, runai = ["runai-model-streamer[s3,gcs,azure]>=0.17.0"], in the style of the mx extra. The 0.17.0 floor is the release that adds the tensor_names and owned arguments of SafetensorsStreamer. The s3, gcs and azure extras install the readers for s3://, gs:// and az://. The README and the guide tell users to pip install "tensorrt_llm[runai]". This is a new optional dependency, so it needs the license and vulnerability scan from the checklist.

Limits: Kimi K3 only, the checkpoint needs an index file, and no draft model or LoRA on object storage.

Test Coverage

  • tests/unittest/llmapi/test_object_storage.py (new, added to the l0_cpu list): URI detection, metadata pull, completion marker, concurrent pulls, force pull.
  • tests/unittest/llmapi/test_llm_args.py and test_llm_utils.py: policy validation, draft model and LoRA rules, metadata pull on local rank 0.
  • tests/unittest/_torch/pyexecutor/test_model_loader_checkpoint_io.py: the URI and metadata directory reach Model Streamer.

Not run yet on this rebased branch: the unit tests above (the changed Python files compile). The rebase onto the current #20782 head had conflicts in l0_cpu.yml and the Kimi K3 README, resolved by hand. CI will exercise the tests.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

yijingl-nvidia and others added 13 commits November 3, 2026 15:02
@yijingl-nvidia
Refactor checkpoint I/O policy to a string enum
8308e59 
Use a new str enum CheckpointIOPolicy for better code readability and maintenance.
Public API and telemetry are updated accordingly.

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
@yijingl-nvidia
Refactor built-in checkpoint formats to a string enum
d35e37c 
Use a new str enum CheckpointFormat for built-in checkpoint formats like "HF", "MX" and "mistral".
Custom 3rd-party formats are still accepted as str.
This improves code readability and maintenance.

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
@yijingl-nvidia
Remove dead code; Refactor checkpoint loading helpers to take LoadFormat
df74402 
tighten function parameter load_format type range from LoadFormat | str to LoadFormat.
This is dead code as the upstream functions don't use checkpoint loading helper functions with str load formats.

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
@yijingl-nvidia
Refactor checkpoint I/O policy resolution
4240151 
Refactor the function that generates checkpoint io policy fallback/error reason
so it's easier for humans to read.
Also rename _construct_checkpoint_loader to construct_checkpoint_loader since it's
used outside the file defining it.

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
@yijingl-nvidia
Refactor ModelLoader.load() and related code
0c49500 
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>
@yijingl-nvidia
remove MX only interface from BaseCheckpointLoader
fc912b0 
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
@yijingl-nvidia
Move Kimi K3 weight loading into modeling_kimi_k3_loader.py
a5425df 
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>
@yijingl-nvidia
Add Model Streamer checkpoint loading
3f9998b 
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>
@yijingl-nvidia
Add selective Model Streamer loading for Kimi K3
db0aa91 
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>
@benamizilberNvidia
Add object storage helpers for model metadata
2ec77af 
Add llmapi/object_storage.py for loading a model from an s3://, gs:// or
az:// URI. It detects object-storage URIs and copies the files that Hugging
Face loaders read from local disk (config, tokenizer, remote code, not the
weights) into a node-local cache directory named after a hash of the URI.

A .runai_complete marker makes later pulls no-ops, and concurrent processes on
a node wait for the first one. TRTLLM_OBJECT_STORAGE_FORCE_PULL=1 downloads
the files again once per launch.

Add the unit tests to the CPU test list.

Signed-off-by: Ben Ami Zilber <bzilber@nvidia.com>
@benamizilberNvidia
Accept object-storage model URIs in the LLM API
ede32d2 
An object-storage model needs checkpoint_io_policy="model_streamer"; any other
policy raises before anything is downloaded. A draft model or LoRA directory
on object storage is rejected, a local one is allowed.

The config, generation config and tokenizer are read from the local metadata
copy, which this change pulls once per node:
- llm_args: the tokenizer validator loads a given tokenizer from the copy.
- llm_utils: CachedModelLoader pulls the metadata on the local rank 0 of each
  node.
- llm: the HF config, the generation config and the default tokenizer use the
  local copy instead of the URI.

Signed-off-by: Ben Ami Zilber <bzilber@nvidia.com>
@benamizilberNvidia
Stream Kimi K3 weights from object storage
de6d56d 
ModelLoader hands the object-storage URI, and the local metadata directory,
to Model Streamer instead of the local checkpoint directory. The streamer's
tensor source map reads model.safetensors.index.json from the local copy and
maps each tensor to its file URI in the bucket.

A separate draft checkpoint is still read from a local path with the
checkpoint loader.

Signed-off-by: Ben Ami Zilber <bzilber@nvidia.com>
@benamizilberNvidia
Document loading Kimi K3 from object storage
fa1b9ac 
Describe the URI forms, the required model_streamer policy, the metadata
cache and its force-pull option, and the limits (Kimi K3 only, an index file
is required, no draft model or LoRA on object storage) in the Kimi K3 README
and the checkpoint loading guide.

Signed-off-by: Ben Ami Zilber <bzilber@nvidia.com>
@benamizilberNvidia benamizilberNvidia changed the title [None][feat] Load Kimi K3 weights from object storage with Model Streamer [feat] Load Kimi K3 weights from object storage with Model Streamer Nov 4, 2026
@benamizilberNvidia
Add a runai extra for loading from object storage
74e864d 
Declare the runai-model-streamer dependency as an optional extra, as the mx
extra does for modelexpress: runai-model-streamer[s3,gcs,azure]>=0.17.0. The
0.17.0 floor is the release that adds the tensor_names and owned arguments of
SafetensorsStreamer; the extras install the readers for s3://, gs:// and az://.

Tell users to install it in the checkpoint loading guide and the Kimi K3
README.

Signed-off-by: Ben Ami Zilber <bzilber@nvidia.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@benamizilberNvidia benamizilberNvidia

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@benamizilberNvidia @yijingl-nvidia