Skip to content

fix(inference): skip generation VAE for reasoner-only runs#132

Open
chrisvoncsefalvay wants to merge 2 commits into
NVIDIA:mainfrom
chrisvoncsefalvay:feat/reasoner-conditional-tokenizer
Open

fix(inference): skip generation VAE for reasoner-only runs#132
chrisvoncsefalvay wants to merge 2 commits into
NVIDIA:mainfrom
chrisvoncsefalvay:feat/reasoner-conditional-tokenizer

Conversation

@chrisvoncsefalvay

Copy link
Copy Markdown

Reasoner-only CLI inference to no longer instantiate the generation vision tokeniser or resolves the video VAE. Mixed and generative input sets retain the existing eager-loading behaviour.

Root cause

OmniMoTModel.set_up_tokenizers instantiated model.config.tokenizer unconditionally. For Cosmos3-Edge, that tokeniser resolves Wan2.2_VAE.pth, even when inference only produces reasoner text.

Changes

  • detect all-reasoner input sets before model construction and apply a narrow model.config.load_vision_tokenizer=false override
  • make the generation vision tokeniser optional only under that override; the default remains enabled
  • use configured compression factors when no live tokeniser exists, covering network construction and token-budget batching
  • treat separate vision-decode GPU placement as a no-op when the generation tokenizer is intentionally absent
  • cover all-reasoner, mixed, empty, default-load, skip-load, batching, and separate-decode cases

Behaviour

  • Reasoner-only inputs: the VLM processor and checkpoint-local vision tower still load; the generation VAE does not.
  • Mixed or generative inputs: unchanged; the generation tokenizer follows the existing eager path.
  • Other entry points: unchanged by default because load_vision_tokenizer defaults to true.

Fixes #128

Detect all-reasoner CLI batches before model construction and disable only the generation vision tokenizer. Keep mixed and generative workloads on the existing eager path.
Use configured compression factors for token-budget batching when the generation tokenizer is absent, and treat separate vision decode GPU placement as a no-op in that state.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Cosmos3-Edge reasoner-only inference initialises the 2.8 GB video VAE

1 participant