Implement decoder-only GPT-style transformer

Goal

Implement a simplified decoder-only Transformer language model (similar in spirit to GPT) for next-token prediction. The implementation should be modular, using four main classes:

1. MultiHeadAttention
2. FeedForward
3. DecoderLayer
4. GPT (the full model)

You may assume a deep learning framework (e.g., PyTorch or TensorFlow), but you should clearly specify tensor shapes and operations.

Model details and requirements

Assume:

• Batch size: B
• Sequence length: T
• Embedding dimension: d_model
• Number of attention heads: num_heads (assume d_model is divisible by num_heads )
• Vocabulary size: V

1. Multi-head self-attention (MultiHeadAttention)

Implement a class MultiHeadAttention that performs masked self-attention:

• Inputs:
  • x : Tensor of shape (B, T, d_model) (token embeddings or layer outputs).
• Internally learnable parameters:
  • Linear projections to compute queries Q , keys K , and values V for all heads.
• Operations:
  1. Project x into Q , K , and V with shapes (B, T, d_model) .
  2. Split them into num_heads heads, each of dimension d_head = d_model / num_heads .
  3. Compute scaled dot-product attention for each head:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_{head}}} + M\right)V,$

 where `M` is a **causal mask** that prevents a position from attending to future positions (positions `> t`).

4. Concatenate the outputs of all heads back to shape (B, T, d_model). 5. Apply a final linear projection to return to (B, T, d_model).

• Output:
  • Tensor of shape (B, T, d_model) .

Ensure you correctly implement the causal (autoregressive) mask so that position t can only attend to positions 0..t.

2. Position-wise feed-forward network (FeedForward)

Implement a class FeedForward for the position-wise MLP:

• Inputs:
  • x : Tensor of shape (B, T, d_model) .
• Internal structure (typical choice):
  • Linear layer from d_model to d_ff (e.g., 4 * d_model ).
  • Non-linear activation (e.g., GELU or ReLU).
  • Linear layer from d_ff back to d_model .
• All operations are applied independently to each position in the sequence.
• Output:
  • Tensor of shape (B, T, d_model) .

3. Decoder layer (DecoderLayer)

Implement a single Transformer decoder block that combines attention, feed-forward, residual connections, and layer normalization:

• Inputs:
  • x : Tensor of shape (B, T, d_model) .
• Components:
  1. LayerNorm ln1 before self-attention.
  2. MultiHeadAttention block (masked self-attention).
  3. Residual connection: x = x + attention_output .
  4. LayerNorm ln2 before feed-forward.
  5. FeedForward block.
  6. Residual connection: x = x + ff_output .
• Output:
  • Tensor of shape (B, T, d_model) .

Your DecoderLayer should be reusable so that multiple layers can be stacked.

4. GPT model (GPT)

Implement a GPT class representing the full decoder-only Transformer:

• Inputs:
  • input_ids : Integer tensor of shape (B, T) representing token indices.
• Components:
  1. Token embedding layer mapping token IDs to vectors of size d_model .
  2. Positional encodings (learned or fixed) of shape (T, d_model) added to the token embeddings.
  3. A stack of N identical DecoderLayer blocks.
  4. A final linear layer mapping from d_model to vocabulary size V .
• Forward pass:
  1. Embed tokens, add positional encodings to obtain (B, T, d_model) .
  2. Pass through the N decoder layers sequentially.
  3. Apply the final linear layer to obtain logits of shape (B, T, V) .
• Output:
  • Logits for next-token prediction: (B, T, V) .

Additional notes

• You should define the forward method for each class with correct tensor transformations.
• Be careful about tensor shapes when splitting/combining heads.
• Ensure the causal mask is correctly broadcast and applied in attention.
• You do not need to implement training loops or optimization; focus on correct and clean model implementation.