DOCSEARCHbyAlgolia
Submit Docs

CALM-pytorch (LLM Augmented LLMs)

repository·main·Indexed 41 days ago

https://github.com/lucidrains/calm-pytorch

A PyTorch implementation of Google DeepMind research for composing an 'anchor' LLM with multiple 'augmentation' LLMs, including Vision Transformers, via cross-attention. It provides the CALM class for model composition, a FineTuner wrapper for training with 🤗 Accelerate, and AugmentParams for configuring connectivity and multi-modal inputs.

Tokens
1.4K
Snippets
5
Records
5
Agent score
2%

What's inside CALM-pytorch

  1. Use CALM for fine-tuning and generation

    main

    Fine-tuning

    Call the calm instance with seq, mask, and prompt to compute the loss. The prompt can be a single Tensor or a tuple of Tensors if using multiple modalities (e.g., text and vision).

    Generation

    Use the .generate() method to prompt the composed model. You must provide the prompt and the desired seq_len.

    # Fine-tuning forward pass
loss = calm(
    seq,
    mask=mask,
    prompt=prompt
)
loss.backward()

# Generation
generated = calm.generate(
    prompt=seq[:, :1],
    seq_len=1024
)

  3. Configure CALM with Vision Transformers (ViT)

    main

    To use a Vision Transformer as an augmentation model, use AugmentParams with input_shape, hidden_position, and an extract_blocks_fn. The extract_blocks_fn is used to identify the specific modules (like Attention blocks) that will be used for connectivity.

    Example configuration for a ViT:

    AugmentParams(
    model=vit,
    input_shape=(3, 256, 256),
    hidden_position='input',
    extract_blocks_fn=lambda vit: [m for m in vit.modules() if isinstance(m, Attention)]
)
    from vit_pytorch.vit import ViT, Attention
from CALM_pytorch import CALM, AugmentParams

# ... setup anchor_llm, augment_llm1, augment_llm2, vit ...

calm = CALM(
    anchor_llm=anchor_llm,
    augment_llms=(
        AugmentParams(model=augment_llm1, mask_kwarg='mask'),
        AugmentParams(model=augment_llm2, mask_kwarg='mask'),
        AugmentParams(
            model=vit,
            input_shape=(3, 256, 256),
            hidden_position='input',
            extract_blocks_fn=lambda vit: [m for m in vit.modules() if isinstance(m, Attention)]
        )
    ),
    attn_kwargs=dict(
        linear_project_context=True,
        pre_rmsnorm=True,
        flash=True
    )
)

# Prompt can be a tuple for multi-modal inputs
prompt = (
    torch.randint(0, 20000, (1, 256)),
    torch.randint(0, 20000, (1, 256)),
    torch.randn(1, 3, 256, 256)
)

loss = calm(seq, mask=mask, prompt=prompt)

  4. Train CALM using the FineTuner class

    main

    The FineTuner class provides a convenient wrapper around 🤗 Accelerate for training.

    Dataset Requirements: The dataset argument should return a dictionary of input kwargs to calm (e.g., dict(seq=Tensor, mask=Tensor, prompt=Tensor)). Alternatively, it can return a Tuple, in which case you must specify the correct ordered value of kwarg names via data_kwargs (though the provided example uses the dictionary approach).

    from CALM_pytorch import FineTuner

trainer = FineTuner(
    calm=calm,
    dataset=dataset,
    batch_size=16,
    num_train_steps=10000,
    learning_rate=3e-4,
    weight_decay=1e-2,
    warmup_steps=1000,
    checkpoint_every=1000
)

trainer()

  5. Initialize the CALM model

    main

    The CALM class implements the LLM Augmented LLMs architecture. It requires an anchor_llm and one or more augment_llms wrapped in AugmentParams.

    AugmentParams allows you to specify the augmentation model and its connectivity settings, such as connect_every_num_layers or custom connections (tuples of anchor-to-augment layer indices).

    from CALM_pytorch import CALM, AugmentParams

# Basic usage with one augmentation LLM
calm = CALM(
    anchor_llm=anchor_llm,
    augment_llms=AugmentParams(
        model=augment_llm,
        connect_every_num_layers=4
    )
)

# Usage with multiple augmentation LLMs and custom connectivity
calm = CALM(
    anchor_llm=anchor_llm,
    augment_llms=(
        AugmentParams(
            model=augment_llm1,
            connections=((1, 12), (2, 12)), # (augment_layer, anchor_layer)
        ),
        AugmentParams(
            model=augment_llm2,
            connections=((6, 1), (12, 12)),
        )
    )
)