AutoModel/Nemotron Omni truncates excess image features when rollout placeholders disagree
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Problem
The Nemotron Omni AutoModel implementation used by NeMo-RL catches an image-feature/placeholder shape mismatch, logs a warning, and continues by slicing the projected vision features to the number of image placeholders. This can produce training inputs whose visual conditioning differs from rollout generation while retaining the rollout's token IDs and log probabilities.
For example, if the rollout contains 100 `` positions but learner-side preprocessing produces pixels that encode to 400 projected features, the sequence still has 100 image positions. The fallback inserts the first 100 of the 400 features. That is not equivalent to encoding the smaller image used during generation.
## Source references
Verified against public source on September 15, 2026:
- NeMo-RL main `c69a64274b34cf31f53722ab0cd46ec0637ef273` [invokes the image processor with image-only dummy text and attaches its media tensors without replacing rollout token IDs](https://github.com/NVIDIA-NeMo/RL/blob/c69a64274b34cf31f53722ab0cd46ec0637ef273/nemo_rl/data/multimodal_utils.py#L1455-L1487). This helper does not reconcile the projected image-feature count with the rollout's placeholder count.
- That RL revision [pins AutoModel to `1814c6c93a66b9d59d254960ef6a99a64249b671`](https://github.com/NVIDIA-NeMo/RL/tree/c69a64274b34cf31f53722ab0cd46ec0637ef273/3rdparty/Automodel-workspace/Automodel).
- The pinned AutoModel's [dynamic-resolution insertion path catches the mismatch and truncates `vit_embeds`](https://github.com/NVIDIA-NeMo/Automodel/blob/1814c6c93a66b9d59d254960ef6a99a64249b671/nemo_automodel/components/models/nemotron_omni/model.py#L915-L926):
```python
try:
inputs_embeds[selected] = inputs_embeds[selected] * 0.0 + vit_embeds
except Exception as e:
logger.warning(...)
n_token = int(selected.sum().item())
inputs_embeds[selected] = inputs_embeds[selected] * 0.0 + vit_embeds[:n_token]
```
The same fallback remains on [AutoModel main `7fe64e71fadd60096b6fb393d0158e3f6926759a`](https://github.com/NVIDIA-NeMo/Automodel/blob/7fe64e71fadd60096b6fb393d0158e3f6926759a/nemo_automodel/components/models/nemotron_omni/model.py#L1118-L1129). The tile-based image branch also has a truncating fallback.
## Minimal component reproduction
This isolates the feature-insertion operation, without requiring model weights or GPUs:
```python
import torch
inputs_embeds = torch.zeros(100, 4) # 100 authoritative rollout placeholders
selected = torch.ones(100, dtype=torch.bool)
vit_embeds = torch.arange(400 * 4, dtype=torch.float32).reshape(400, 4)
try:
inputs_embeds[selected] = inputs_embeds[selected] * 0.0 + vit_embeds
except Exception as e:
print(f"Shape mismatch: {e}")
n_token = int(selected.sum().item())
inputs_embeds[selected] = inputs_embeds[selected] * 0.0 + vit_embeds[:n_token]
assert torch.equal(inputs_embeds, vit_embeds[:100])
print("Insertion returned normally after dropping 300 vision features")
```
I also executed the exact dynamic-resolution `try`/`except` extracted from both linked AutoModel revisions with CPU tensors (PyTorch 2.11.0). Both returned normally for 100 placeholders / 400 features. With two images, 100 placeholders per image and 400 features per image, the flattened fallback takes the first 200 features, all belonging to image 1; image 2's positions therefore receive image 1's features.
## Expected behavior
- Do not treat prefix truncation as a repair for mismatched visual conditioning. Report a clear alignment error instead of continuing with the wrong features.
- Ensure the NeMo-RL rollout-to-learner path reproduces or transports the generator's actual image representation: image identity/order, resize geometry and processing settings, or processed pixels. Preserve the authoritative rollout IDs.
- Validate alignment per image, not only the aggregate count. Equal token counts alone also do not prove that geometry or pixels match.
- Add regression coverage for excess features, insufficient features, and multiple images, including a matching-count control.
## Validation scope
The truncating fallback and the NeMo-RL attachment behavior are source-verified; the insertion behavior is reproduced with the exact source block. A generator using a smaller effective image allocation than the learner is one possible trigger, not a claim that this occurs in every configuration. No end-to-end RL run, workload incidence, log-probability impact, or TMPE attribution is claimed by this report.
Filing here to track the NeMo-RL AutoModel integration; the truncation itself is implemented in the AutoModel dependency.
Contributor guide
Assessment
This issue has not been assessed yet.