菜单
Train large models efficiently across heterogeneous hardware (GPUs, TPUs, Trainium) using strict encapsulation principles, achieving constant code complexity when adding features across hundreds of modules.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill modular-large-model-training命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
正在显示 SKILL.md
基于 SOC 职业分类
| name | modular-large-model-training |
| title | AXLearn: Modular Large Model Training on Heterogeneous Infrastructure |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2507.05411 |
| keywords | ["Distributed Training","Model Training Framework","Hardware Agnosticism","JAX","Modularity"] |
| description | Train large models efficiently across heterogeneous hardware (GPUs, TPUs, Trainium) using strict encapsulation principles, achieving constant code complexity when adding features across hundreds of modules. |
Training large models requires coordinating across thousands of code modules—embeddings, attention, normalization, optimizers, checkpointing, monitoring. Adding a single feature (e.g., Rotary Position Embeddings) traditionally requires modifying O(N) modules where N scales with system complexity. This technical debt makes research and deployment difficult. Strict encapsulation—forcing modules to communicate only through well-defined interfaces—enables constant O(1) complexity feature addition regardless of system size.
AXLearn demonstrates that composition-based design eliminates this scaling problem, enabling the same architecture to train efficiently on GPUs, TPUs, and AWS Trainium without framework changes. The system achieves competitive inference performance by reusing training components, reducing the overhead of maintaining separate inference stacks.
Most ML frameworks use subtyping and inheritance, where features propagate through class hierarchies causing O(N) or O(N²) changes. AXLearn enforces strict encapsulation: components don't inherit from common base classes but communicate through explicit interfaces. This inverts complexity—adding features requires only changing code that uses that feature, not everything that could use it.
The key insight is that functional programming (JAX) + strict encapsulation enables both modularity and hardware agnosticism simultaneously. Rather than hardware-specific code scattered throughout, a single compiler layer (XLA) handles hardware details, allowing identical module code to run on different accelerators.
The following implements modular architecture principles for scalable model training.
Step 1: Strict Encapsulation Module Design
This demonstrates how to design self-contained modules with encapsulation.
import jax
import jax.numpy as jnp
from typing import NamedTuple, Any, Callable
class ModuleOutput(NamedTuple):
"""Explicit interface: modules only produce outputs of known type."""
value: jnp.ndarray
state: dict # Module internal state changes
class StrictModule:
"""Base pattern for strictly encapsulated modules."""
def __init__(self, config: dict):
"""Modules are configured at construction, not modified later."""
self.config = config
self.params = None
def initialize_parameters(self, key: jax.random.PRNGKey, input_shape: tuple):
"""Explicit parameter initialization with shape contract."""
raise NotImplementedError
def __call__(
self,
inputs: jnp.ndarray,
training: bool = True,
context: dict = None
) -> ModuleOutput:
"""
Single call signature regardless of module type.
No hidden dependencies on module type or inheritance.
"""
raise NotImplementedError
class LinearLayer(StrictModule):
"""Example: strictly encapsulated linear layer."""
def __init__(self, config: dict):
super().__init__(config)
self.output_dim = config.get("output_dim", 768)
self.use_bias = config.get("use_bias", True)
def initialize_parameters(self, key, input_shape):
"""Deterministic initialization based on input shape."""
input_dim = input_shape[-1]
key_w, key_b = jax.random.split(key)
# He initialization
self.params = {
"weight": jax.random.normal(key_w, (input_dim, self.output_dim)) * jnp.sqrt(2.0 / input_dim),
}
if self.use_bias:
self.params["bias"] = jnp.zeros(self.output_dim)
def __call__(self, inputs: jnp.ndarray, training: bool = True, context: dict = None) -> ModuleOutput:
"""Standard linear transformation with no side effects."""
output = jnp.dot(inputs, self.params["weight"])
if self.use_bias:
output = output + self.params["bias"]
return ModuleOutput(
value=output,
state={} # No state changes
)
class AttentionLayer(StrictModule):
"""Example: strictly encapsulated attention."""
def __init__(self, config: dict):
super().__init__(config)
self.num_heads = config.get("num_heads", 12)
self.head_dim = config.get("hidden_dim", 768) // self.num_heads
def initialize_parameters(self, key, input_shape):
hidden_dim = input_shape[-1]
key_q, key_k, key_v, key_out = jax.random.split(key, 4)
scale = jnp.sqrt(2.0 / hidden_dim)
self.params = {
"q_proj": jax.random.normal(key_q, (hidden_dim, hidden_dim)) * scale,
"k_proj": jax.random.normal(key_k, (hidden_dim, hidden_dim)) * scale,
"v_proj": jax.random.normal(key_v, (hidden_dim, hidden_dim)) * scale,
"out_proj": jax.random.normal(key_out, (hidden_dim, hidden_dim)) * scale,
}
def __call__(self, inputs: jnp.ndarray, training: bool = True, context: dict = None) -> ModuleOutput:
batch, seq_len, hidden_dim = inputs.shape
# Project
Q = jnp.dot(inputs, self.params["q_proj"])
K = jnp.dot(inputs, self.params["k_proj"])
V = jnp.dot(inputs, self.params["v_proj"])
# Reshape for multi-head
Q = Q.reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
K = K.reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
V = V.reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
# Attention
scores = jnp.matmul(Q, K.transpose(0, 1, 3, 2)) / jnp.sqrt(self.head_dim)
attn_weights = jax.nn.softmax(scores, axis=-1)
attn_output = jnp.matmul(attn_weights, V)
# Reshape back
attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch, seq_len, hidden_dim)
# Output projection
output = jnp.dot(attn_output, self.params["out_proj"])
return ModuleOutput(
value=output,
state={}
)
Step 2: Composable Module Hierarchy
This shows how modules compose without breaking encapsulation.
class TransformerBlock(StrictModule):
"""Composing modules: attention + FFN block."""
def __init__(self, config: dict):
super().__init__(config)
self.attention = AttentionLayer(config)
self.ffn_layer1 = LinearLayer({"output_dim": config.get("ffn_dim", 3072)})
self.ffn_layer2 = LinearLayer({"output_dim": config.get("hidden_dim", 768)})
self.norm1 = config.get("norm_fn", "layer_norm")
def initialize_parameters(self, key, input_shape):
key_attn, key_ffn1, key_ffn2 = jax.random.split(key, 3)
self.attention.initialize_parameters(key_attn, input_shape)
self.ffn_layer1.initialize_parameters(key_ffn1, input_shape)
ffn_output_shape = input_shape[:-1] + (self.config.get("ffn_dim", 3072),)
self.ffn_layer2.initialize_parameters(key_ffn2, ffn_output_shape)
# Params from submodules
self.params = {
"attention": self.attention.params,
"ffn": {
"layer1": self.ffn_layer1.params,
"layer2": self.ffn_layer2.params,
}
}
def __call__(self, inputs: jnp.ndarray, training: bool = True, context: dict = None) -> ModuleOutput:
# Attention with residual
self.attention.params = self.params["attention"]
attn_out = self.attention(inputs, training, context)
x = inputs + attn_out.value
# Layer norm + FFN with residual
x_norm = jax.nn.layer_norm(x)
self.ffn_layer1.params = self.params["ffn"]["layer1"]
ffn_out1 = self.ffn_layer1(x_norm, training, context)
ffn_out1_activated = jax.nn.gelu(ffn_out1.value)
self.ffn_layer2.params = self.params["ffn"]["layer2"]
ffn_out2 = self.ffn_layer2(ffn_out1_activated, training, context)
x = x + ffn_out2.value
return ModuleOutput(
value=x,
state={**attn_out.state, **ffn_out2.state}
)
class TransformerModel(StrictModule):
"""Full transformer: composition of blocks."""
def __init__(self, config: dict):
super().__init__(config)
self.num_layers = config.get("num_layers", 12)
self.blocks = [TransformerBlock(config) for _ in range(self.num_layers)]
def initialize_parameters(self, key, input_shape):
keys = jax.random.split(key, self.num_layers)
self.params = {}
current_shape = input_shape
for i, block in enumerate(self.blocks):
block.initialize_parameters(keys[i], current_shape)
self.params[f"block_{i}"] = block.params
def __call__(self, inputs: jnp.ndarray, training: bool = True, context: dict = None) -> ModuleOutput:
x = inputs
all_states = {}
for i, block in enumerate(self.blocks):
block.params = self.params[f"block_{i}"]
output = block(x, training, context)
x = output.value
all_states.update(output.state)
return ModuleOutput(
value=x,
state=all_states
)
Step 3: Hardware-Agnostic Training Loop
This implements training that works across GPUs, TPUs, and other hardware without modification.
import optax
class HardwareAgnosticTrainer:
"""Training loop transparent to hardware backend."""
def __init__(self, model: StrictModule, config: dict):
self.model = model
self.config = config
self.optimizer = optax.adam(learning_rate=config.get("learning_rate", 1e-3))
@jax.jit # JIT compiles to any backend automatically
def train_step(self, params: dict, batch: dict, opt_state: dict):
"""Single training step (JIT-compiled for any hardware)."""
def loss_fn(params):
self.model.params = params
output = self.model(batch["input_ids"], training=True)
logits = output.value
# Compute cross-entropy loss
targets = batch["labels"]
loss = optax.softmax_cross_entropy_with_integer_labels(logits, targets).mean()
return loss
loss, grads = jax.value_and_grad(loss_fn)(params)
# Update parameters
updates, opt_state = self.optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)
return params, opt_state, loss
def train_epoch(self, train_loader, num_epochs: int = 3):
"""Training epoch (runs on whatever hardware JAX detects)."""
# Initialize
key = jax.random.PRNGKey(0)
dummy_input = jnp.ones((1, 512, 768)) # Batch of sequences
self.model.initialize_parameters(key, dummy_input.shape)
opt_state = self.optimizer.init(self.model.params)
for epoch in range(num_epochs):
for batch in train_loader:
self.model.params, opt_state, loss = self.train_step(
self.model.params, batch, opt_state
)
print(f"Epoch {epoch + 1}/{num_epochs} | Loss: {loss:.4f}")
Step 4: Feature Addition Complexity Analysis
This demonstrates O(1) feature addition complexity.
class LoC_ComplexityAnalyzer:
"""Measure lines of code required to add features."""
@staticmethod
def analyze_feature_addition(feature_name: str, module_count: int) -> dict:
"""
Hypothetical feature addition: Add RotaryPositionalEmbedding to attention layers.
In traditional frameworks (subtyping): O(N) where N = module count
In AXLearn (strict encapsulation): O(1)
"""
# AXLearn approach: only modify AttentionLayer
axlearn_changes = {
"files_modified": 1, # Only attention.py
"new_classes": 1, # RotaryEmbedding
"modified_methods": 1, # AttentionLayer.__call__
"total_loc_added": 50, # ~50 lines
"complexity": "O(1)"
}
# Traditional approach: modify every attention usage
traditional_changes = {
"files_modified": module_count, # Every attention use site
"modified_methods": module_count,
"total_loc_added": module_count * 10, # ~10 lines per site
"complexity": f"O(N) where N={module_count}"
}
return {
"axlearn": axlearn_changes,
"traditional": traditional_changes,
"complexity_improvement": module_count
}
Hyperparameters and Configuration
| Parameter | Recommended Value | Range | Notes |
|---|---|---|---|
| Batch Size | 32-128 per device | 1-1024 | Adjust for memory constraints of hardware |
| Learning Rate | 1e-3 | 1e-5 to 1e-1 | Use warmup + decay schedule |
| Gradient Accumulation | 1-4 steps | 1-16 | Simulate larger batches on memory-limited hardware |
| Checkpoint Frequency | Every 1000 steps | 100-10000 | More frequent on unstable hardware |
| Max Sequence Length | 512-2048 | 128-16384 | Hardware-dependent; TPU handles longer better |
| Mixed Precision (bfloat16) | Enabled | True/False | Reduces memory 2×, better on TPU than GPU |
When to Use
When NOT to Use
Common Pitfalls
AXLearn: Modular Large Model Training on Heterogeneous Infrastructure. https://arxiv.org/abs/2507.05411