feat(models): add Inkling VLM MoE support - #3095
Conversation
|
/claude review |
|
/ok to test 7fab517 |
|
Reviewed the Inkling VLM MoE onboarding (model/layers/adapter/processing, registry, collate + PP media, MoE parallelizer fp32 handling, PP functional metadata, grad-norm warmup, recipe/docs). Accounted for all 26 changed files; Strong points worth calling out: registry entry + Findings are minor (3 inline): tensor-contract docstrings are missing/incomplete on new tensor-accepting functions with nonstandard interleaved/flat layouts ( No blocking issues found. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: hemildesai <hemild@nvidia.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
7fab517 to
15724fb
Compare
Signed-off-by: hemildesai <hemild@nvidia.com>
Signed-off-by: hemildesai <hemild@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
|
Hi @hemildesai Thanks for the great work. Since #2998, my earlier TieSupport refactor, has merged, every model with an LM head now needs to declare its tying policy. I added a small follow-up marking Inkling as UNTIED_ONLY with focused tests. This matches its existing separate embed/unembed weights and should not change any of your implementation’s behavior. |
|
/ok to test b4336ee |
|
/claude review |
|
/ok to test 1dafb36 |
|
Light review complete. This is a well-constructed model onboarding PR (Inkling multimodal MoE VLM) plus a Torch 2.12 pipelining compatibility path. Coverage is strong: state-dict round-trip, forward/backward gradient parity, multimodal tower parity, two-stage pipeline forward parity, and the static-metadata warm-up path all have focused CPU tests. One minor finding posted inline:
No correctness, security, or gradient-handling issues found. |
| if stage.group_size > 2: | ||
| edges.add((0, stage.group_size - 1)) | ||
|
|
||
| edge_phases: list[list[tuple[int, int]]] = [] |
There was a problem hiding this comment.
This new production helper accepts inputs_meta / outputs_meta tensor tuples but its docstring doesn't document their layout, which the repo's tensor-contract invariant requires for functions accepting tensor inputs. Please add a Google-style Args section, e.g.:
Args:
stage: The pipeline stage to configure.
inputs_meta: Meta-device tensors describing the stage's forward inputs —
`[microbatch, sequence]` int64 for the first stage, else
`[microbatch, sequence, hidden]` in the model dtype.
outputs_meta: Meta-device tensors describing the stage's forward outputs —
`[microbatch, sequence, vocab]` for a stage holding the LM head, else
`[microbatch, sequence, hidden]`.
With transformers pinned at main's 5.12.1 (no Inkling upstream), the Inkling unit tests are importorskip-guarded and skip in CI, so the new nemo_automodel/components/models/inkling code has ~0% patch coverage and trips the default 80% codecov/patch gate (10.82% of diff hit). Follow the existing Bagel precedent: exclude the inkling package from the default patch group and add an informational-only inkling group. Coverage activates automatically once a transformers release shipping Inkling lets the guarded tests run. Signed-off-by: hemildesai <hemild@nvidia.com>
|
/claude review |
|
/ok to test 5172c25 |
|
Light review complete — 19 changed files reviewed (uv.lock excluded). This PR onboards the Inkling multimodal MoE VLM (model wrapper, EP MoE layers, state-dict adapter, processor helper), extends the VLM collate/PP-media path, and adds Torch 2.12 pipeline-stage static-metadata + neighbor-warmup compatibility. Overall the implementation looks solid and unusually well tested:
One minor in-scope finding posted inline: the overridden No correctness, security, or config-boundary issues found. |
| past_key_values: Any | None = None, | ||
| inputs_embeds: torch.FloatTensor | None = None, | ||
| use_cache: bool | None = None, | ||
| **kwargs: Any, |
There was a problem hiding this comment.
This overriding forward has no docstring, yet it materially changes tensor handling versus the HF base: input_ids here may carry either token ids or float hidden states (the pipeline-stage path sets inputs_embeds = input_ids when input_ids.dtype.is_floating_point), which is a non-obvious dual contract. Per the repo's tensor-contract rule, please document the layouts of the tensor inputs, e.g.:
Args:
input_ids: Tensor of shape [batch, sequence] of token ids on the
first stage; on later pipeline stages without an embedding table
it instead carries float hidden states of shape
[batch, sequence, hidden].
inputs_embeds: Optional Tensor of shape [batch, sequence, hidden].
attention_mask: Tensor of shape [batch, sequence] or a precomputed
per-attention-type mask mapping.
The public forward at line 278 (pixel_values, input_ids, inputs_embeds, ...) similarly documents no input layouts and would benefit from the same Args section.
|
/claude review |
|
/ok to test d754b04 |
@hemildesai, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
|
Review summary — reviewed all 19 changed files (excluding uv.lock). This is a well-structured onboarding of the Inkling multimodal MoE VLM: the model wrapper reuses HF towers, adds an EP InklingMoE, a state-dict adapter with an exact round-trip, plus a Torch 2.12 PP static-metadata compatibility path and generalized chunk_step3_media flat-pixel handling. Test coverage is strong and behavior-asserting: CPU state-dict round-trip, forward-logit parity (KL < 1e-3), backward gradient parity across gate/experts/embeddings, multimodal tower parity, and a two-stage PP forward-parity test. ModelCapabilities correctly declares supports_pp/supports_ep, RoPE stays HF-owned, and cast_model_to_dtype is used with strict fp32 holders for the router bias and short convolutions — so the low-precision-RoPE and checkpoint-free-init concerns are handled. Findings (2, non-blocking): both are tensor-contract docstring gaps on the two new public forward overrides in inkling/model.py (InklingTextModel.forward and InklingForConditionalGeneration.forward), posted inline. These new overrides accept/return tensors whose layouts — especially the PP-stage-dependent input/output reinterpretation — the repo requires be documented, and they diverge from the base contract so cannot simply reference it. No correctness, security, gradient-domain, or API-surface issues found. The num_patches addition in default_collate_fn is appropriately guarded (image_token_id present + 5-D pixel_values) so it does not perturb existing Qwen/Gemma paths. |
| inputs_embeds: torch.FloatTensor | None = None, | ||
| use_cache: bool | None = None, | ||
| **kwargs: Any, | ||
| ) -> BaseModelOutputWithPast: |
There was a problem hiding this comment.
InklingTextModel.forward is a new override that materially changes the base contract (it reinterprets input_ids as pre-computed hidden states on non-embedding PP stages, per line 87-90), so it can't defer to HF's documented contract and needs its own tensor-layout docstring. Per the repo's tensor-contract rule, add a Google-style Args/Returns documenting at least:
Args:
input_ids: Tensor of shape [batch, sequence] of token ids on the first
stage; on later PP stages a float hidden-states tensor of shape
[batch, sequence, hidden] passed in this slot (input_ids is then set
to None internally).
inputs_embeds: Optional Tensor of shape [batch, sequence, hidden].
attention_mask: Tensor of shape [batch, sequence] or a mapping of
per-attention-type masks.
Returns:
BaseModelOutputWithPast whose last_hidden_state has shape
[batch, sequence, hidden].
| logits_to_keep: int | torch.Tensor = 0, | ||
| **kwargs: Any, | ||
| ) -> Any: | ||
| """Run the standard Inkling forward or its pipeline-stage equivalent.""" |
There was a problem hiding this comment.
This new public forward takes several tensor inputs (input_ids, pixel_values, inputs_embeds, audio_input_ids, ...) and returns either hidden states [batch, sequence, hidden] (non-last PP stage) or logits [batch, sequence, unpadded_vocab] (last stage), but the one-line docstring documents none of these layouts. Per the repo's tensor-contract rule, add a Google-style Args/Returns giving each tensor input's shape (e.g. input_ids: Tensor of shape [batch, sequence]; pixel_values: Tensor of shape [num_images, temporal_patch, patch, patch, channels]) and describing the two possible return layouts, since callers on different PP stages consume different outputs.
3 participants
Summary
torch_mmexperts_fp32_paramsholders that reuse the existing FSDP holder path;components/moe/parallelizer.pyis unchanged frommainValidation
14060480: 256 H100 GPUs, PP8/EP32, global batch 256, sequence length 20480:0using NCCL's default communicator blocking behavior2.3000to1.7906; final grad norm0.7824, memory58.53 GiB, and throughput4,415.29tokens/s1.7898, grad norm0.8266, memory58.54 GiB, and throughput4,316.08tokens/s14060405, exit0:0)14060410, exit0:0)mainMoE parallelizer (Slurm job14057832, exit0:0)uv lock --check,compileall, andgit diff --checkCaveat