案例：用 TorchTitan 训自定义架构

TorchTitan 默认只支持 Llama 3 / DeepSeek-V3 / Mixtral 等几个模型。本章手把手加一个新模型（假设叫 MyArch）。

12.1需要写哪些文件

12.2model.py 骨架

# 标准 PyTorch nn.Module，不要任何分布式
class MyTransformerBlock(nn.Module):
    def __init__(self, cfg: MyArchArgs):
        super().__init__()
        self.attention = Attention(cfg)
        self.feed_forward = FeedForward(cfg)
        self.attention_norm = RMSNorm(cfg.dim)
        self.ffn_norm = RMSNorm(cfg.dim)
    def forward(self, x, freqs_cis):
        h = x + self.attention(self.attention_norm(x), freqs_cis)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

class MyModel(nn.Module):
    def __init__(self, cfg: MyArchArgs):
        super().__init__()
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.layers = nn.ModuleDict({str(i): MyTransformerBlock(cfg) for i in range(cfg.n_layers)})
        self.norm = RMSNorm(cfg.dim)
        self.output = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
    def forward(self, ids):
        h = self.embed(ids)
        for layer in self.layers.values():
            h = layer(h, freqs_cis=...)
        return self.output(self.norm(h))

注意：写的是普通 PyTorch nn.Module，不要在这里搞 ParallelLinear。所有并行在 parallelize.py 注入。

12.3args.py

@dataclass
class MyArchArgs:
    dim: int = 2048
    n_layers: int = 16
    n_heads: int = 16
    vocab_size: int = 32000
    max_seq_len: int = 4096
    ...
    # 多种 flavor
    @classmethod
    def from_flavor(cls, flavor):
        return {
            "1B": cls(dim=2048, n_layers=16, n_heads=16),
            "7B": cls(dim=4096, n_layers=32, n_heads=32),
        }[flavor]

12.4parallelize.py

from torch.distributed.tensor.parallel import (
    ColwiseParallel, RowwiseParallel, parallelize_module
)
from torch.distributed.fsdp import fully_shard

def parallelize_my_arch(model, mesh, parallel_dims, config):
    # TP
    if parallel_dims.tp_enabled:
        tp_mesh = mesh["tp"]
        for layer in model.layers.values():
            plan = {
                "attention.wq": ColwiseParallel(),
                "attention.wk": ColwiseParallel(),
                "attention.wv": ColwiseParallel(),
                "attention.wo": RowwiseParallel(),
                "feed_forward.w1": ColwiseParallel(),
                "feed_forward.w2": RowwiseParallel(),
                "feed_forward.w3": ColwiseParallel(),
            }
            parallelize_module(layer, tp_mesh, plan)
        # embed / output 也要 TP
        parallelize_module(model.embed, tp_mesh, RowwiseParallel())
        parallelize_module(model.output, tp_mesh, ColwiseParallel())

    # AC
    if config.activation_checkpoint.mode != "none":
        for layer in model.layers.values():
            apply_ac_to_layer(layer, config.activation_checkpoint)

    # FSDP2
    if parallel_dims.dp_shard_enabled:
        dp_mesh = mesh["dp_shard"]
        for layer in model.layers.values():
            fully_shard(layer, mesh=dp_mesh)
        fully_shard(model, mesh=dp_mesh)

    return model

12.5__init__.py：注册

from .model import MyModel
from .args import MyArchArgs
from .parallelize import parallelize_my_arch

def build_my_model(args: MyArchArgs) -> MyModel:
    return MyModel(args)

# 在 torchtitan/models/__init__.py 里注册
from torchtitan.models.my_arch import build_my_model, MyArchArgs, parallelize_my_arch

MODELS_REGISTRY["my_arch"] = {
    "build_model_fn": build_my_model,
    "args_cls": MyArchArgs,
    "parallelize_fn": parallelize_my_arch,
}

12.6train_configs/my_arch.toml

[model]
name = "my_arch"
flavor = "7B"

[training]
batch_size = 1
seq_len = 4096
tensor_parallel_degree = 2
data_parallel_shard_degree = 4
...

12.7跑

torchrun --nproc_per_node=8 torchtitan/train.py \
    --job.config_file ./train_configs/my_arch.toml

能跑通就 OK。后续可以加新的 flavor、调 TP plan、加 PP 切层。

12.8注意事项

• weight 命名要跟 HF 一致（方便 ckpt 转）；
• 用 RMSNorm 不用 LayerNorm（性能差不少）；
• Attention 用 SDPA（torch 2.5+ 自动选 flash 实现）；
• layers 用 ModuleDict 而非 ModuleList（FSDP2 prefetch 更友好）。

12.9这章你需要带走的

• 加新架构 = 4 个文件 + 1 个 toml；
• model 写普通 nn.Module，并行在 parallelize 注入；
• TP plan 关键是 QKV 列切 + O 行切；MLP gate/up 列切 + down 行切；
• 看 torchtitan/models/llama/ 是最好的模板。