Codebase Structure

class Mlp(nn.Module):
    def __init__(self, in_features: int,
            hidden_features=None, out_features=None,
            bias: bool = False):
        super().__init__()
        out_features    = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x

class Attention(nn.Module):
    def __init__(self, dim: int, head_dim: int = 64,
            qkv_bias: bool = False, proj_bias: bool = False):
        super().__init__()
        self.num_heads = dim // head_dim
        self.scale     = head_dim ** -0.5

        # Fused QKV projection — one Linear, 3× output dim
        self.qkv          = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_out_proj = nn.Linear(dim, dim, bias=proj_bias)

Attention — Forward & Masking

    def forward(self, x, mask=None):
        B, L, D = x.shape

        # Fused QKV → reshape → (3, B, H, L, d_k)
        qkv = self.qkv(x).reshape(B, L, 3, self.num_heads,
                             D // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)   # each: (B, H, L, d_k)

        attn = (q @ k.transpose(-2, -1)) * self.scale  # (B, H, L, L)
        if mask is not None:
            mask = rearrange(mask, 'b n m -> b 1 n m')
            attn = attn.masked_fill(~mask,
                           -torch.finfo(attn.dtype).max)
        attn = attn.softmax(dim=-1)
        x = (attn @ v).transpose(1, 2).reshape(B, L, D)
        return self.attn_out_proj(x)

Block — Pre-Norm + Residuals

class Block(nn.Module):
    def __init__(self, dim: int, head_dim: int = 64,
            mlp_ratio: float = 4.0, use_bias: bool = False):
        super().__init__()
        self.norm1 = LayerNorm(dim)
        self.attn  = Attention(dim, head_dim=head_dim,
                       qkv_bias=use_bias, proj_bias=use_bias)
        self.norm2 = LayerNorm(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp   = Mlp(in_features=dim,
                      hidden_features=mlp_hidden_dim, bias=use_bias)

    def forward(self, x, mask=None):
        x = x + self.attn(self.norm1(x), mask)  # Xₐ = X + Attn(LN₁(X))
        x = x + self.mlp(self.norm2(x))         # Xᵦ = Xₐ + MLP(LN₂(Xₐ))
        return x

GPT — forward_model

    def forward_model(self, x):
        B, L = x.size()

        # Input + positional embedding
        x = self.input_embedding(x) + self.positional_embedding[:L]

        # Causal mask: tril(-inf matrix).bool() → True=attend, False=masked
        causal_mask = torch.tril(torch.full(
            (L, L), float('-inf'), device=x.device), diagonal=0).bool()
        causal_mask = causal_mask.unsqueeze(0)  # (1, L, L)

        x = self.trunk(x, mask=causal_mask)
        return self.to_logits(self.out_norm(x))  # (B, L, vocab_size)

GPT — compute_ce_loss & forward

    def compute_ce_loss(self, logits, target_seq, padding_idx=-100):
        B, L, vocab_size = logits.shape
        return F.cross_entropy(
            logits.reshape(-1, vocab_size),
            target_seq.reshape(-1),
            ignore_index=padding_idx
        )

    def forward(self, data_dict):
        seq = data_dict[self.seq_read_key]  # (B, L+1)
        input_seq  = seq[:, :-1]           # [SOS, t1, …, tL] — drop last
        target_seq = seq[:, 1:]            # [t1, …, EOS] — drop first (shift!)
        logits = self.forward_model(input_seq)
        loss   = self.compute_ce_loss(logits, target_seq, self.padding_idx)
        return loss, {'ppl': torch.exp(loss)}

GPT — generate

    @torch.no_grad()
    def generate(self, context=[0], eos_idx=None,
            temp=1.0, top_p=0.0, top_k=0.0):
        was_training = self.training
        self.eval()

        # Initialize the sequence with the start-of-sequence token
        current_tokens = torch.tensor(
            [context], dtype=torch.long, device=self.device)

        for _ in range(self.max_seq_len - len(context)):
            logits = self.forward_model(current_tokens)
            # Get logits for the last token and reshape to (1, vocab_size)
            next_logits = logits[0, -1, :].unsqueeze(0)
            next_token, _ = sample_tokens(next_logits,
                temperature=temp, top_k=top_k, top_p=top_p)
            current_tokens = torch.cat(
                [current_tokens, next_token.unsqueeze(1)], dim=1)
            if next_token.item() == eos_idx:
                break

        if was_training:
            self.train()

        return current_tokens

Q2.1 — Effect of Temperature

# sample_tokens applies temp like:
logits = logits / temp
probs  = F.softmax(logits, dim=-1)
token  = torch.multinomial(probs, 1)

Q2.2 — top_k and top_p