For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

VocabParallelEmbedding

VocabParallelEmbedding

class max.nn.VocabParallelEmbedding(vocab_size, hidden_dim, dtype, devices, quantization_encoding=None, name=None)

source

Bases: Module

A lookup table for embedding integer indices into dense vectors.

This layer works like Embedding except the embedding table is sharded on the vocabulary dimension across all devices. When called, VocabParallelEmbedding accepts a TensorValueLike of integer indices along with signal buffers for cross-device communication and returns a list of TensorValue tensors (one per device) containing the corresponding embedding vectors.

from max.driver import Accelerator, accelerator_count
from max.dtype import DType
from max.graph import DeviceRef
from max.nn import VocabParallelEmbedding

# builds an Allreduce, so it requires at least one GPU.
if accelerator_count() > 0:
    device_ref = DeviceRef.from_device(Accelerator())
    embedding_layer = VocabParallelEmbedding(
        vocab_size=1000,
        hidden_dim=256,
        dtype=DType.float32,
        devices=[device_ref],
        name="embeddings",
    )

Initializes the vocab-parallel embedding layer.

Parameters:

• vocab_size (int) – The number of unique items in the vocabulary. Indices must be in the range [0, vocab_size).
• hidden_dim (int) – The dimensionality of each embedding vector.
• dtype (DType) – The data type of the embedding weights.
• devices (list[DeviceRef]) – The devices where embedding lookups are executed. Model init transfers the initially CPU-resident weights to this device.
• quantization_encoding (QuantizationEncoding | None) – Optional quantization encoding for the weights.
• name (str | None) – The name identifier for the embedding weight matrix.