Nemotron-3.5-Lightning-1.4B-A0.1B-MTP
This is a tiny version of nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4 created for testing and development. MTP (Multi-Token Prediction) tensors are included in the checkpoint for testing MTP quantization pipelines.
Model Details
- Base Model: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4
- Architecture: nemotron_h (hybrid Mamba + MoE)
- Total Parameters: 1.4B (main: 1.0B + MTP layers: 0.4B)
- Activated Parameters: ~0.1B (MoE with 2/16 active experts per token)
- Weight dtype: bfloat16 (unquantized float; base model is NVFP4)
Configuration Changes
| Parameter | Original | Tiny |
|---|---|---|
num_hidden_layers | 52 | 12 |
hidden_size | 2688 | 2048 |
n_routed_experts | 128 | 16 |
num_attention_heads | 32 | 16 |
mamba_num_heads | 64 | 32 |
intermediate_size | 1856 | 1024 |
moe_intermediate_size | 1856 | 1024 |
moe_shared_expert_intermediate_size | 3712 | 2048 |
vocab_size | 131072 | 131072 (unchanged) |
num_nextn_predict_layers | 1 | 1 (unchanged) |
mtp_layers_block_type | ["full_attention", "moe"] | ["full_attention", "moe"] (unchanged) |
MTP Layer Details
The checkpoint includes 69 MTP tensors (~405M params) in a separate shard model_mtp.safetensors. These are registered in model.safetensors.index.json under the mtp.* prefix. The MTP block mirrors the base model's mtp_layers_block_type: ["full_attention", "moe"] with 16 routed experts.
Note: The HuggingFace NemotronHForCausalLM class silently ignores mtp.* keys at load time (_keys_to_ignore_on_load_unexpected = [r"mtp.*"]), so MTP tensors must be loaded separately for quantization testing.
Checkpoint Structure
The checkpoint is sharded into two files:
model.safetensors— main model weights (247 tensors, ~4.4GB)model_mtp.safetensors— MTP layer weights (69 tensors, ~0.4GB)model.safetensors.index.json— unified weight map covering both shards
Usage
from transformers import AutoModelForCausalLM, AutoTokenizer
# Main model loads without MTP (HF class ignores mtp.* keys)
model = AutoModelForCausalLM.from_pretrained("inference-optimization/Nemotron-3.5-Lightning-1.4B-A0.1B-MTP")
tokenizer = AutoTokenizer.from_pretrained("inference-optimization/Nemotron-3.5-Lightning-1.4B-A0.1B-MTP")
input_ids = tokenizer("According to all known laws", return_tensors="pt").input_ids.to(model.device)
output = model.generate(input_ids, max_new_tokens=20)
print(tokenizer.decode(output[0]))
Validation Output
Loading weights: 100%|██████████| 97/97
Success: 1.0000896453857422 <= 10.0
Generating sample text:
According to all known laws of aviation, there is no way a bee should be able to fly.
Notes
- This model is NOT intended for inference. It is a synthetic tiny model for testing MTP quantization pipelines in llm-compressor.
- Weights are random bfloat16 floats, fine-tuned on a toy copypasta dataset to verify the training loop works.
- The base model is NVFP4 quantized; this tiny model uses full float weights.
- MTP tensors are present in the checkpoint but ignored by the standard HF loader — this is intentional, matching the behavior of the full-size model.
Creation Process
Created using the llm-compressor create-tiny-model Claude skill. The architecture was validated by instantiating on meta device, fine-tuned on a toy dataset to perplexity ~1.0, and MTP tensors were added as synthetic random bfloat16 weights matching the mtp_layers_block_type configuration.