AR / ar /model /model.py

Upload AR-4B pretraining codebase, configs, and cluster scripts

2ae6691 verified 18 days ago

3.5 kB

import torch
import torch.nn as nn

from .config import ARConfig
from .rmsnorm import RMSNorm
from .block import DecoderBlock
from .rope import build_rope_cache


class ARModel(nn.Module):
    """Decoder-only transformer: token IDs in, next-token scores out."""

    def __init__(self, cfg: ARConfig):
        super().__init__()
        self.cfg = cfg

        # token embedding: maps each token ID to a dim-wide vector
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)

        # the stack of identical decoder blocks
        self.blocks = nn.ModuleList([
            DecoderBlock(cfg.dim, cfg.n_heads, cfg.n_kv_heads,
                         cfg.head_dim, cfg.hidden_dim, cfg.rms_norm_eps)
            for _ in range(cfg.n_layers)
        ])

        self.final_norm = RMSNorm(cfg.dim, cfg.rms_norm_eps)
        # output head: dim-wide vector -> score for every token in the vocab
        self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)

        # tie weights: embedding table and output head share one matrix
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight

        # precompute RoPE tables once; reused by every block (not a parameter)
        cos, sin = build_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        # initialize all weights small; without this, 24 stacked layers make
        # the output logits explode (std ~32 instead of ~1)
        self.apply(self._init_weights)
        # extra: scale down the two residual-path projections per block by
        # 1/sqrt(2*n_layers) so the signal doesn't compound across depth
        scale = (2 * cfg.n_layers) ** -0.5
        for block in self.blocks:
            torch.nn.init.normal_(block.attn.o_proj.weight, mean=0.0, std=0.02 * scale)
            torch.nn.init.normal_(block.mlp.down_proj.weight, mean=0.0, std=0.02 * scale)

    def _init_weights(self, module):
        # small normal init for all linear and embedding weights
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, token_ids: torch.Tensor, block_mask) -> torch.Tensor:
        B, T = token_ids.shape

        # look up embeddings, then slice the RoPE tables to this sequence length
        x = self.embed(token_ids)
        cos, sin = self.rope_cos[:T], self.rope_sin[:T]

        # pass through every decoder block in turn
        for block in self.blocks:
            x = block(x, cos, sin, block_mask)

        # final clean-up, then project to vocab-sized scores (logits)
        x = self.final_norm(x)
        return self.lm_head(x)



    # def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
    #     B, T = token_ids.shape

    #     # look up embeddings, then slice the RoPE tables to this sequence length
    #     x = self.embed(token_ids)
    #     cos, sin = self.rope_cos[:T], self.rope_sin[:T]

    #     # pass through every decoder block in turn
    #     for block in self.blocks:
    #         x = block(x, cos, sin)

    #     # final clean-up, then project to vocab-sized scores (logits)
    #     x = self.final_norm(x)
    #     return self.lm_head(x)

	import torch
	import torch.nn as nn

	from .config import ARConfig
	from .rmsnorm import RMSNorm
	from .block import DecoderBlock
	from .rope import build_rope_cache


	class ARModel(nn.Module):
	"""Decoder-only transformer: token IDs in, next-token scores out."""

	def __init__(self, cfg: ARConfig):
	super().__init__()
	self.cfg = cfg

	# token embedding: maps each token ID to a dim-wide vector
	self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)

	# the stack of identical decoder blocks
	self.blocks = nn.ModuleList([
	DecoderBlock(cfg.dim, cfg.n_heads, cfg.n_kv_heads,
	cfg.head_dim, cfg.hidden_dim, cfg.rms_norm_eps)
	for _ in range(cfg.n_layers)
	])

	self.final_norm = RMSNorm(cfg.dim, cfg.rms_norm_eps)
	# output head: dim-wide vector -> score for every token in the vocab
	self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)

	# tie weights: embedding table and output head share one matrix
	if cfg.tie_embeddings:
	self.lm_head.weight = self.embed.weight

	# precompute RoPE tables once; reused by every block (not a parameter)
	cos, sin = build_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
	self.register_buffer("rope_cos", cos, persistent=False)
	self.register_buffer("rope_sin", sin, persistent=False)

	# initialize all weights small; without this, 24 stacked layers make
	# the output logits explode (std ~32 instead of ~1)
	self.apply(self._init_weights)
	# extra: scale down the two residual-path projections per block by
	# 1/sqrt(2*n_layers) so the signal doesn't compound across depth
	scale = (2 * cfg.n_layers) ** -0.5
	for block in self.blocks:
	torch.nn.init.normal_(block.attn.o_proj.weight, mean=0.0, std=0.02 * scale)
	torch.nn.init.normal_(block.mlp.down_proj.weight, mean=0.0, std=0.02 * scale)

	def _init_weights(self, module):
	# small normal init for all linear and embedding weights
	if isinstance(module, nn.Linear):
	torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
	if module.bias is not None:
	torch.nn.init.zeros_(module.bias)
	elif isinstance(module, nn.Embedding):
	torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

	def forward(self, token_ids: torch.Tensor, block_mask) -> torch.Tensor:
	B, T = token_ids.shape

	# look up embeddings, then slice the RoPE tables to this sequence length
	x = self.embed(token_ids)
	cos, sin = self.rope_cos[:T], self.rope_sin[:T]

	# pass through every decoder block in turn
	for block in self.blocks:
	x = block(x, cos, sin, block_mask)

	# final clean-up, then project to vocab-sized scores (logits)
	x = self.final_norm(x)
	return self.lm_head(x)



	# def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
	# B, T = token_ids.shape

	# # look up embeddings, then slice the RoPE tables to this sequence length
	# x = self.embed(token_ids)
	# cos, sin = self.rope_cos[:T], self.rope_sin[:T]

	# # pass through every decoder block in turn
	# for block in self.blocks:
	# x = block(x, cos, sin)

	# # final clean-up, then project to vocab-sized scores (logits)
	# x = self.final_norm(x)
	# return self.lm_head(x)