nGPT-pytorch

Install the nGPT-pytorch package using pip to access the normalized GPT implementation.

$ pip install nGPT-pytorch

Instead of standard additive residuals, nGPT uses a Residual module that implements a Slerp-like (Spherical Linear Interpolation) update combined with L2 normalization.

This mechanism uses a Scale module to learn the interpolation factor. The update follows this logic:

1. Compute the output of the sub-module (fn).
2. Apply L2 normalization to the output.
3. Perform a linear interpolation (lerp) between the original input x and the normalized output, scaled by a learned branch_scale.
4. Apply L2 normalization to the resulting interpolated value.

This approach helps maintain the norm of activations across deep networks.

The architecture relies on two specialized components to maintain normalization:

1. Scale: A module that maintains a learnable scale parameter. It returns self.scale * self.forward_scale, where forward_scale is initialized as init / scale. This allows the model to control the magnitude of activations or residuals while keeping the effective initialization controlled.

2. NormLinear: A linear layer that applies L2 normalization to its weights.

   • If parametrize=True, it uses PyTorch's register_parametrization to ensure the weights are always L2-normalized.
   • It supports groups (hyperspheres), allowing the weights to be normalized in chunks.
   • The output is scaled by self.scale (calculated as groups ** -1).

The nGPT class implements a normalized GPT architecture that learns on the hypersphere. You can initialize the model with specific dimensions and depth, and optionally enable attention normalization for Query and Key vectors.

To perform training, call the model with return_loss = True. To perform inference or get logits, call the model without that flag.

import torch
from nGPT_pytorch import nGPT

model = nGPT(
    num_tokens = 256,
    dim = 512,
    depth = 4,
    attn_norm_qk = True
)

x = torch.randint(0, 256, (2, 2048))

# Training mode: returns the loss
loss = model(x, return_loss = True)
loss.backward()

# Inference mode: returns logits of shape (batch, sequence, num_tokens)
logits = model(x) # (2, 2048, 256)

The nTransformer allows for per-layer control of scaling factors. If you provide a single value, it is applied to all layers. If you provide a tuple, it must match the depth of the transformer.

Supported scaling parameters:

• alpha_init: Sets the initial scale for all residual connections (defaults to 1. / depth).
• alpha_attn_init / alpha_attn_scale: Scaling for attention residual connections.
• alpha_ff_init / alpha_ff_scale: Scaling for feedforward residual connections.
• s_qk_init / s_qk_scale: Scaling for query/key products in attention.
• s_ff_hidden_init / s_ff_hidden_scale: Scaling for feedforward hidden states.
• s_ff_gate_init / s_ff_gate_scale: Scaling for feedforward gate states.
• s_logit_init / s_logit_scale: Scaling for logits (if applicable).

The nGPT constructor accepts several hyperparameter groups to control the effective learning rates and scales throughout the network. Most of these accept either a single float (applied to all layers) or a tuple[float, ...] of length depth (providing a specific value per layer).

Residual Scaling (Alpha):

• alpha_init: Base initialization for all residuals.
• alpha_attn_init / alpha_attn_scale: Scaling for attention residuals.
• alpha_ff_init / alpha_ff_scale: Scaling for feedforward residuals.

Internal Component Scaling (S):

• s_qk_init / s_qk_scale: Scaling for Query-Key projections.
• s_ff_hidden_init / s_ff_hidden_scale: Scaling for FeedForward hidden dimension.
• s_ff_gate_init / s_ff_gate_scale: Scaling for FeedForward gate dimension.
• s_logit_init / s_logit_scale: Scaling for the output logits.

The cross_over_feedforward function implements a breeding strategy for FeedForward layers. It splits the dim_inner of the layer into two halves using randomly permuted indices and applies crossover to all internal components: to_hidden, to_gate, to_out, hidden_scale, and gate_scale.

Parameters:

• parent1: The first FeedForward parent.
• parent2: The second FeedForward parent.
• child (optional): An existing FeedForward instance to populate. If None, a copy of parent1 is used.

from nGPT_pytorch.evo import cross_over_feedforward

# Assuming parent1 and parent2 are FeedForward instances
child = cross_over_feedforward(parent1, parent2)

The NormLinear class provides a linear layer where weights are constrained by L2 normalization.

It supports two modes:

1. Parametrized: Uses torch.nn.utils.parametrize to ensure the weight remains normalized during optimization.
2. Manual: Manually applies L2 normalization to weights.

It also supports groups (via the num_hyperspheres concept) to allow normalization across specific sub-dimensions.

from nGPT_pytorch.nGPTExperimental import NormLinear

# A linear layer with parametrized L2-normalized weights
layer = NormLinear(dim = 768, dim_out = 768, parametrize = True)

The nGPT_pytorch package provides access to the nGPT model, the nTransformer architecture, and several core building blocks used in the normalized GPT implementation. You can import these directly from the root package.

from nGPT_pytorch import (
    nGPT,
    nTransformer,
    Attention,
    FeedForward,
    NormLinear,
    Residual,
    Scale,
    l2norm
)