adu2021 Deploy Efficient Reasoning Models via Nested Submodel Elasticity
Training multiple model sizes (12B, 9B, 6B) separately requires training three independent models—expensive and wasteful. Nemotron Elastic embeds multiple nested submodels within a single parent model, sharing weights hierarchically. An end-to-end trained router selects which submodel to extract at inference based on computational budget. This enables deploying various model sizes with constant memory overhead while training only once.
The approach achieves 360× cost reduction vs. training family members separately, enabling efficient deployment of reasoning models across diverse hardware constraints.
Core Concept
Model families (like GPT-3 with 7B, 13B, 175B variants) typically train separately, multiplying training cost. Nemotron Elastic inverts this: train once, extract multiple sizes.
The key insight is weight sharing with hierarchical structure:
- Parent Model: Largest model (12B); contains full capability
- Nested Submodels: Extractable smaller models (9B, 6B, 3B) using prefix of parent's weights
- Router: Learnable component selecting which submodel to use based on task/budget
- Elastification Techniques: Domain-specific tricks for Mamba (SSM), MLPs, and layer selection
This enables:
- Single training pass (using curriculum on budget diversity)
- Multiple deployable sizes
- Constant memory (all sizes fit in parent's memory)
- 360× cost reduction vs. training separate models
Architecture Overview
- Nested Model Structure: Submodel i uses first k_i layers, sharing weights with parent; i.e., all weights are reused
- Router-Based Selection: Learned module predicting optimal submodel for task/input; can be based on task type, input complexity, or deployment budget
- Elastification Techniques: Specialized methods for Mamba (group-aware SSM elastification), MLPs (heterogeneous elastification), and layer selection (normalized MSE-based importance)
- Two-Stage Training Curriculum: Stage 1: train parent + routers; Stage 2: distillation from parent to submodels
- Knowledge Distillation: Multi-budget optimization ensuring each submodel learns from parent
Implementation Steps
Step 1: Nested Model Architecture.
import torch
import torch.nn as nn
class NestedElasticModel(nn.Module):
"""
Parent model with embedded nested submodels.
Each submodel can be extracted without additional training.
"""
def __init__(
self,
vocab_size=32000,
hidden_dim=4096,
num_layers=36,
num_heads=32,
moe_mode='mamba', # 'mamba' or 'transformer'
submodel_depths=[24, 18, 12] # Extract at these layer counts
):
super().__init__()
self.hidden_dim = hidden_dim
self.num_layers = num_layers
self.submodel_depths = submodel_depths
# Shared embedding
self.embedding = nn.Embedding(vocab_size, hidden_dim)
# All layers (shared by all submodels)
if moe_mode == 'mamba':
self.layers = nn.ModuleList([
MambaBlock(hidden_dim) for _ in range(num_layers)
])
else: # Transformer
self.layers = nn.ModuleList([
TransformerBlock(hidden_dim, num_heads) for _ in range(num_layers)
])
# Output head
self.output_proj = nn.Linear(hidden_dim, vocab_size)
# Router: selects which depth to use
self.router = Router(hidden_dim, len(submodel_depths))
def forward(self, x, output_router_logits=False):
"""
Forward through full parent model.
"""
hidden = self.embedding(x)
for layer in self.layers:
hidden = layer(hidden)
logits = self.output_proj(hidden)
if output_router_logits:
router_logits = self.router(hidden)
return logits, router_logits
return logits
def extract_submodel(self, submodel_idx):
"""
Extract nested submodel of specified size.
Returns inference-only module.
"""
depth = self.submodel_depths[submodel_idx]
class SubModel(nn.Module):
def __init__(self, parent, depth):
super().__init__()
self.embedding = parent.embedding
self.layers = parent.layers[:depth]
self.output_proj = parent.output_proj
def forward(self, x):
hidden = self.embedding(x)
for layer in self.layers:
hidden = layer(hidden)
return self.output_proj(hidden)
return SubModel(self, depth)
Step 2: Router Module—Dynamic Submodel Selection.
class Router(nn.Module):
"""
Learns to select optimal submodel for given input/task.
Can be conditioned on task type, input complexity, deployment budget.
"""
def __init__(self, hidden_dim, num_submodels):
super().__init__()
self.router_head = nn.Sequential(
nn.Linear(hidden_dim, 256),
nn.ReLU(),
nn.Linear(256, num_submodels)
)
def forward(self, hidden_states):
"""
Predict submodel logits.
hidden_states: (batch, seq_len, hidden_dim) or (batch, hidden_dim)
"""
if len(hidden_states.shape) == 3:
# Aggregate over sequence
hidden = hidden_states.mean(dim=1) # (batch, hidden_dim)
else:
hidden = hidden_states
router_logits = self.router_head(hidden) # (batch, num_submodels)
return router_logits
def select_submodel(self, router_logits):
"""Select submodel greedily (argmax)."""
return torch.argmax(router_logits, dim=-1)
def sample_submodel(self, router_logits, temperature=1.0):
"""Sample submodel (for training diversity)."""
probs = torch.softmax(router_logits / temperature, dim=-1)
return torch.multinomial(probs, 1).squeeze(-1)
Step 3: Elastification Techniques—Depth-Aware Layer Selection.
class LayerImportanceEstimator(nn.Module):
"""
Estimate which layers are most important for each submodel depth.
Use normalized MSE-based layer importance.
"""
def __init__(self, num_layers):
super().__init__()
# Learn importance scores for each layer
self.layer_importance = nn.Parameter(torch.ones(num_layers))
def compute_importance(self):
"""Normalize to probability distribution."""
return torch.softmax(self.layer_importance, dim=0)
def select_important_layers(self, num_select):
"""Select top-k important layers."""
importance = self.compute_importance()
_, indices = torch.topk(importance, num_select)
return sorted(indices.tolist())
def elastify_depth(parent_model, submodel_depth, parent_depth):
"""
Elastification: create submodel using important layers.
For simplicity, use prefix of layers (can extend to importance-based selection).
"""
class ElasticSubmodel(nn.Module):
def __init__(self, parent, depth):
super().__init__()
self.embedding = parent.embedding
self.layers = parent.layers[:depth]
self.output_proj = parent.output_proj
def forward(self, x):
hidden = self.embedding(x)
for layer in self.layers:
hidden = layer(hidden)
return self.output_proj(hidden)
return ElasticSubmodel(parent_model, submodel_depth)
class MambaElastification(nn.Module):
"""
Elastification for Mamba SSM models: group-aware SSM elastification.
"""
def __init__(self, hidden_dim, num_groups=4):
super().__init__()
self.num_groups = num_groups
self.group_projections = nn.ModuleList([
nn.Linear(hidden_dim, hidden_dim // num_groups) for _ in range(num_groups)
])
def elastify_ssm(self, ssm_layer, target_hidden_dim):
"""
Reduce SSM hidden dimension while maintaining group structure.
"""
class ElasticSSM(nn.Module):
def __init__(self, parent_ssm, target_dim, group_proj):
super().__init__()
self.parent_ssm = parent_ssm
self.target_dim = target_dim
self.group_proj = group_proj
def forward(self, x):
# Reduce via learned projection
x_reduced = self.group_proj(x)
# Forward through parent SSM (handles partial dims)
# Note: real implementation would adapt SSM internals
output = self.parent_ssm(x_reduced)
return output
return ElasticSSM(ssm_layer, target_hidden_dim, self.group_projections[0])
Step 4: Two-Stage Training Curriculum.
def train_elastic_model(
model, train_loader, num_epochs=100,
submodel_depths=[24, 18, 12], lr=5e-5
):
"""
Two-stage training: Stage 1 trains parent + routers; Stage 2 distillation.
"""
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
# Stage 1: Train parent + router (50 epochs)
print("Stage 1: Training parent model and router...")
for epoch in range(num_epochs // 2):
total_loss = 0
for batch in train_loader:
input_ids = batch['input_ids']
target_ids = batch['target_ids']
budget_level = batch.get('budget_level', torch.zeros(input_ids.shape[0])) # Optional
# Forward: get parent logits and router logits
parent_logits, router_logits = model(input_ids, output_router_logits=True)
# Loss 1: Parent model loss
parent_loss = torch.nn.functional.cross_entropy(
parent_logits.reshape(-1, parent_logits.shape[-1]),
target_ids.reshape(-1)
)
# Loss 2: Router loss (encourage selecting appropriate submodel)
# Router logits should align with complexity of input
router_target = estimate_task_difficulty(input_ids)
router_loss = torch.nn.functional.cross_entropy(router_logits, router_target)
total_loss = parent_loss + 0.1 * router_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
print(f"Epoch {epoch} (Stage 1): loss={total_loss.item():.4f}")
# Stage 2: Knowledge distillation from parent to submodels (50 epochs)
print("\nStage 2: Knowledge distillation to submodels...")
for epoch in range(num_epochs // 2, num_epochs):
total_loss = 0
for batch in train_loader:
input_ids = batch['input_ids']
target_ids = batch['target_ids']
# Parent outputs (frozen or with slow learning)
with torch.no_grad():
parent_logits, _ = model(input_ids, output_router_logits=True)
# Submodel distillation
distill_loss = 0
for submodel_idx, depth in enumerate(submodel_depths):
submodel = model.extract_submodel(submodel_idx)
submodel_logits = submodel(input_ids)
# KL divergence between parent and submodel
kl_loss = torch.nn.functional.kl_div(
torch.log_softmax(submodel_logits, dim=-1),
torch.softmax(parent_logits, dim=-1),
reduction='batchmean'
)
distill_loss += kl_loss / len(submodel_depths)
optimizer.zero_grad()
distill_loss.backward()
optimizer.step()
total_loss += distill_loss
print(f"Epoch {epoch} (Stage 2): distill_loss={distill_loss.item():.4f}")
return model
def estimate_task_difficulty(input_ids):
"""
Heuristic: input complexity determines which submodel.
Can be replaced with learned difficulty estimator.
"""
# Length-based: longer inputs need larger models
lengths = (input_ids > 0).sum(dim=-1)
difficulty = torch.clamp(lengths // 256, 0, 2) # 0-2 for 3 submodels
return difficulty
Practical Guidance
When to Use: Deploying reasoning models across diverse hardware (GPUs, CPUs, edge devices) where model size flexibility is crucial. Cost reduction from single training pass is massive advantage.
Architecture Decisions:
- Submodel depths: geometric progression (e.g., 36→24→12 layers) works well; adjust based on task requirements
- Router mechanism: simple MLPs sufficient; can add task-specific gating for specialized routing
- Elastification: layer prefix works for Transformers; Mamba requires group-aware techniques
- Training schedule: 50/50 split between parent training and distillation; adjust ratio based on dataset size
Pitfalls:
- Router overconfidence: Router may converge to always selecting one submodel; use entropy regularization
- Distillation divergence: Submodels diverging from parent; use stronger KL weight, periodic teacher updates
- Submodel collapse: Smaller submodels may learn trivial solutions; use curriculum learning (start large, gradually shrink)
- Depth selection: Equal-spaced layer selection may be suboptimal; profile on validation set
When NOT to Use: Single-model deployment where multi-size support unnecessary; models with highly specialized architecture (e.g., custom expert selection).
Integration: Works with Transformers and Mamba; minimal changes to base architecture.
Reference: https://arxiv.org/abs/2511.16664