Inference#

After fine-tuning, we run inference to generate predictions and evaluate model quality. Keep the video preprocessing and text prompt style/template consistent with training: frame rate (video_fps) and resolution settings affect how much spatial and temporal detail the model sees, and mismatches in preprocessing or prompting can shift results on fine-grained sports QA.

This section covers distributed inference (interactive and sbatch) for the reference base model used by the current examples and checkpoints from our training stacks. Use the backend that matches how the checkpoint was produced: NeMo AutoModel for Hugging Face weights (base, SFT, or LoRA), or Megatron-Bridge for Megatron-format.

Runbook (commands and smoke examples): avlm/inference/README.md.

Environment variables and overrides: inference environment reference.

Setup#

Inference code lives in the Sports Intelligence Playbooks repo at avlm/inference/ (shared driver: common/run_inference.py). Clone the repo, cd to its root, and run all launch commands from there.

Containers#

Set CONTAINER_IMAGE in each backend’s slurm/launch_local.yaml (copy from launch.yaml).

AutoModel inference: nemo-automodel:26.06
Megatron-Bridge inference: nemo:26.06

On Slurm with enroot, convert the NGC image once to .sqsh (see Setup). Point CONTAINER_IMAGE at the resulting path.

Two Backends (Pick One)#

Pick one of the following based on the framework you used for training.

	NeMo AutoModel (HF)	Megatron-Bridge (Megatron)
Checkpoint	HF base, full SFT weights, or LoRA adapter + base	Bridge DDP (non-FSDP) or FSDP SFT checkpoint, or base checkpoint + LoRA adapter
Container	nemo-automodel:26.06	nemo:26.06
Launcher path	`avlm/inference/automodel/slurm/`	`avlm/inference/megatron-bridge/slurm/`
Example configs	`automodel/configs/default_sft.yaml`, `default_lora.yaml`	`megatron-bridge/configs/default_sft.yaml`, `default_sft_fsdp.yaml`, `default_lora.yaml`, `default_lora_fsdp.yaml`

Shared Layout#

common/run_inference.py — JSONL inference driver (both backends)
common/models/ — backend-specific model loaders
common/utils/_inference_lib.sh — CLI/env overrides for launch scripts
Per-backend configs/ + slurm/ — launch_local.yaml, interactive + sbatch launchers

Outputs: avlm/inference/outputs/<INFERENCE_NAME>/ Logs (batch): avlm/inference/logs/<INFERENCE_NAME>_<timestamp>/

AutoModel Inference#

HF-format inference via common/run_inference.py with model_name: nemotron_omni_hf. Supports the HF base model, full SFT checkpoints, and LoRA adapters.

Full SFT / base: automodel/configs/default_sft.yaml
LoRA: automodel/configs/default_lora.yaml

Set model_path to the base or SFT checkpoint. For LoRA, set model_path to the base model and adapter_path to the adapter checkpoint. See the inference runbook above for interactive, batch, smoke, and override commands.

Megatron-Bridge Inference#

Inference over Megatron-Bridge DDP (non-FSDP) or FSDP checkpoints. SFT checkpoints and base checkpoint + LoRA adapter pairs are supported directly. megatron-bridge/entrypoint.py applies the Bridge AVLM overlay, then calls common/run_inference.py with model_name: nemotron_omni_mbridge.

DDP (non-FSDP): megatron-bridge/configs/default_sft.yaml and default_lora.yaml
FSDP: megatron-bridge/configs/default_sft_fsdp.yaml and default_lora_fsdp.yaml

Set model_path to the base or SFT checkpoint. For LoRA, also set adapter_path to the adapter checkpoint. Optional: MEGATRON_BRIDGE_GIT_BOOTSTRAP=1 when not using /opt/Megatron-Bridge from the container. See the inference runbook above for launch commands.

Inference Config Reference#

Shared Fields (Both Backends)#

Field	Role
`inference_name`	Output subdirectory name under `base_output_dir`; also used as the Slurm job name
`data_path`	Eval JSONL (HF conversation format)
`video_root`	Root directory prepended to relative video paths in each JSONL row
`base_output_dir`	Parent directory for outputs (default `avlm/inference/outputs/`)
`max_new_tokens`	Maximum tokens generated per sample
`video_fps`	Frame sampling rate when decoding video clips
`max_video_frames`	Maximum frames passed to the vision encoder per clip. Sampling first targets `video_fps` (roughly clip duration × fps frames); if that count exceeds this cap, frames are reduced to `max_video_frames` by uniform subsampling (evenly spaced indices across the clip)
`resize`	For AutoModel and Megatron-Bridge, when `true`, preprocess decoded video frames by resizing each frame (as an image) to 512×512 before it is passed to the vision encoder. This is an input-image resize only; it does not change how the encoder itself is configured.

CLI / Environment Overrides#

These are applied on top of the inference YAML (see avlm/inference/README.md for examples).

Override	Role
`INFERENCE_CONFIG`	Path to the inference YAML
`INFERENCE_NAME`	Output subdirectory name (overrides `inference_name` in the YAML)
`DATA_PATH`	Eval JSONL path
`MODEL_PATH`	Checkpoint path (overrides `model_path`)
`ADAPTER_PATH`	LoRA adapter checkpoint path (AutoModel or Megatron-Bridge)
`BASE_OUTPUT_DIR`	Parent output directory
`MAX_INFERENCE_SAMPLES`	Limit inference to the first N JSONL rows (useful for smokes)
`MAX_INFERENCE_SAMPLES_SEED`	Randomly sample the limited rows using this seed; without it, the first N rows are used
`IS_FSDP`	Override `is_fsdp` for Megatron-Bridge inference
`FSDP_SHARDING_STRATEGY`	Override `fsdp_sharding_strategy` for Megatron-Bridge inference
`RESUME=1`	Continue a partially written run from `rank_results/`; use the same config, `INFERENCE_NAME`, and number of ranks as the original job

AutoModel Fields#

Field	Role
`model_name`	Must be `nemotron_omni_hf` for the current reference HF checkpoint
`model_path`	HF model id or local path to base or full SFT weights; for LoRA, this must be the base model
`hf_model_path`	Base HF model used for model code and configuration staging
`adapter_path`	Directory with `adapter_config.json` and `adapter_model.safetensors`; providing it enables LoRA inference

Megatron-Bridge Fields#

Field	Role
`model_name`	Must be `nemotron_omni_mbridge` for the current reference Megatron-Bridge checkpoint
`hf_model_path`	HF model id used for the processor/tokenizer
`model_path`	Base or SFT checkpoint directory in DDP (non-FSDP) or FSDP (`fsdp_dtensor`) format
`adapter_path`	LoRA adapter checkpoint used with the base checkpoint in `model_path`
`is_fsdp`	`true` for an FSDP checkpoint; `false` for DDP (non-FSDP)
`fsdp_sharding_strategy`	FSDP inference sharding strategy: `no_shard` or `optim_grads_params`
`tensor_model_parallel_size`	Tensor parallelism for the inference topology
`expert_model_parallel_size`	Expert parallelism (EP); shards MoE experts across the inference topology
`expert_tensor_parallel_size`	Expert tensor parallelism (ETP); shards each expert’s weights across GPUs, separate from `tensor_model_parallel_size`. Usually `1`.

Pipeline parallelism and context parallelism are fixed at 1 for Megatron-Bridge inference.

LoRA Inference (Adapter-Only Weights)#

LoRA training saves adapter weights only — not a full copy of the 30B base model. Inference therefore always needs the base model plus the adapter, or a merged checkpoint produced offline.

AutoModel#

Use default_lora.yaml with:

model_path — HF base model (for example nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16). Do not point this at a consolidated SFT checkpoint.
adapter_path — training output directory containing adapter_config.json and adapter_model.safetensors (often under .../checkpoints_.../epoch_*_step_* or a model/ subfolder)

At load time the driver reads the LoRA matrices, maps wrapper layer names to the HF model, and adds the low-rank delta directly into the base weight tensors in memory. You do not need to merge adapters to disk before inference, but you do need the base model weights available locally or on HuggingFace Hub.

Megatron-Bridge#

Megatron-Bridge loads LoRA adapters directly with their base checkpoint:

DDP (non-FSDP): use default_lora.yaml with the base checkpoint in model_path, the adapter checkpoint in adapter_path, and is_fsdp: false.
FSDP: use default_lora_fsdp.yaml with the FSDP base checkpoint in model_path, the FSDP adapter checkpoint in adapter_path, and is_fsdp: true.

Merging or exporting the adapter remains optional; see megatron-bridge/lora/LORA_GUIDE.MD for those workflows.

Outputs#

avlm/inference/outputs/<INFERENCE_NAME>/predictions.jsonl
avlm/inference/outputs/<INFERENCE_NAME>/predictions.json

Multi-GPU runs also write per-rank scratch files under rank_results/ until all ranks finish and rank 0 merges.

For MCQ scoring and LLM-judge evaluation of these predictions, see Evaluation.

Quick Start#

Complete Setup in the inference runbook for the selected backend.
Select the matching SFT or LoRA runtime config and set inference_name, data_path, video_root, and model_path; set adapter_path for LoRA. For Megatron-Bridge, select DDP (non-FSDP) or FSDP to match the checkpoint.
Run a small interactive smoke test, then submit a batch run after it passes.

Copy, launch, smoke, batch, resume, and override commands are in the inference runbook.