{"id":15601008,"url":"https://github.com/lucidrains/calm-pytorch","last_synced_at":"2025-10-04T00:55:35.626Z","repository":{"id":216327882,"uuid":"741052407","full_name":"lucidrains/CALM-pytorch","owner":"lucidrains","description":"Implementation of CALM from the paper \"LLM Augmented LLMs: Expanding Capabilities through Composition\", out of Google Deepmind","archived":false,"fork":false,"pushed_at":"2024-09-12T17:26:55.000Z","size":962,"stargazers_count":174,"open_issues_count":3,"forks_count":11,"subscribers_count":8,"default_branch":"main","last_synced_at":"2025-03-30T00:06:16.020Z","etag":null,"topics":["artificial-intelligence","attention-mechanisms","cross-attention","deep-learning","transformers"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-01-09T15:49:20.000Z","updated_at":"2025-03-04T07:25:43.000Z","dependencies_parsed_at":"2024-02-15T18:45:44.403Z","dependency_job_id":"e057a510-7e9a-4ba5-a5bc-1c4ec2fb3d7d","html_url":"https://github.com/lucidrains/CALM-pytorch","commit_stats":{"total_commits":68,"total_committers":1,"mean_commits":68.0,"dds":0.0,"last_synced_commit":"92f8521e476d63c924a2cce40acfc37303a43740"},"previous_names":["lucidrains/calm-pytorch"],"tags_count":44,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCALM-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCALM-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCALM-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FCALM-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/CALM-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247419860,"owners_count":20936012,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","attention-mechanisms","cross-attention","deep-learning","transformers"],"created_at":"2024-10-03T02:11:34.398Z","updated_at":"2025-10-04T00:55:30.599Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"./calm.png\" width=400px/\u003e\n\n## CALM - Pytorch\n\nImplementation of CALM from the paper \u003ca href=\"https://arxiv.org/abs/2401.02412\"\u003eLLM Augmented LLMs: Expanding Capabilities through Composition\u003c/a\u003e, out of Google Deepmind\n\nCan support any number of augmentation LLMs\n\n## Install\n\n```bash\n$ pip install CALM-pytorch\n```\n\n## Appreciation\n\n- \u003ca href=\"https://a16z.com/supporting-the-open-source-ai-community/\"\u003eA16Z Open Source AI Grant Program\u003c/a\u003e and \u003ca href=\"https://huggingface.co/\"\u003e🤗 Huggingface\u003c/a\u003e for the generous sponsorships, as well as my other sponsors, for affording me the independence to open source current artificial intelligence research\n\n## Usage\n\nex. with `x-transformers`\n\n```python\nimport torch\nfrom x_transformers import TransformerWrapper, Decoder\n\naugment_llm = TransformerWrapper(\n    num_tokens = 20000,\n    max_seq_len = 1024,\n    attn_layers = Decoder(\n        dim = 512,\n        depth = 12,\n        heads = 8\n    )\n)\n\nanchor_llm = TransformerWrapper(\n    num_tokens = 20000,\n    max_seq_len = 1024,\n    attn_layers = Decoder(\n        dim = 512,\n        depth = 2,\n        heads = 8\n    )\n)\n\n# import CALM wrapper\n\nfrom CALM_pytorch import CALM, AugmentParams\n\ncalm = CALM(\n    anchor_llm,\n    augment_llms = AugmentParams(\n        model = augment_llm,\n        connect_every_num_layers = 4\n    )\n)\n\n# mock input\n\nseq = torch.randint(0, 20000, (1, 1024))\nmask = torch.ones((1, 1024)).bool()\nprompt = torch.randint(0, 20000, (1, 256))\n\n# forward for finetuning loss\n\nloss = calm(\n    seq,\n    mask = mask,\n    prompt = prompt\n)\n\nloss.backward()\n\n# after much training, prompt the composed model\n\ngenerated = calm.generate(\n    prompt = seq[:, :1],\n    seq_len = 1024\n)\n\n```\n\nTo use a handy trainer class using 🤗 Accelerate, just import `FineTuner` and use as follows\n\n```python\ntrainer = FineTuner(\n    calm = calm,\n    dataset = dataset,   # returns a dictionary of input kwargs to calm - dict(seq: Tensor, mask: Tensor, prompt: Tensor). it can also return a Tuple, in which data_kwargs needs to be set to the correct ordered value of kwarg names\n    batch_size = 16,\n    num_train_steps = 10000,\n    learning_rate = 3e-4,\n    weight_decay = 1e-2,\n    warmup_steps = 1000,\n    checkpoint_every = 1000\n)\n\ntrainer()\n\n# checkpoints of the cross attention parameters will be saved to ./checkpoints every 1000 steps\n```\n\nTo explore multiple augmentation LLMs, simply pass in a list for `augment_llm`\n\nex.\n\n```python\ncalm = CALM(\n    anchor_llm = anchor_llm,\n    augment_llm = [AugmentParams(augment_llm1), AugmentParams(augment_llm2)] # pass in a list of AugmentParams wrapping model and other hparams specific to that transformer\n)\n```\n\nSay you want to explore different types of connectivity between anchor and augmentation model(s), just pass in the connections as a tuple of tuple integer pairs, specifying the anchor to augment layer number.\n\n```python\ncalm = CALM(\n    anchor_llm = anchor_llm,\n    augment_llms = (\n        AugmentParams(\n            model = augment_llm1,\n            connections = (\n                (1, 12),  # 1st layer of augment llm1 attended to by 12th layer of anchor llm\n                (2, 12),\n                (3, 12),\n                (4, 12),\n            ),\n        ),\n        AugmentParams(\n            model = augment_llm2,\n            connections = (\n                (6, 1),   # 6th layer of augment llm2 attended to by 1st layer of anchor llm\n                (6, 2),\n                (12, 12),\n            )\n        )\n    )\n)\n```\n\nCALM setup with 2 specialized augmentation LLMs + a vision transformer\n\n```python\nimport torch\n\n# pip install vit-pytorch x-transformers\n\nfrom vit_pytorch.vit import ViT, Attention\nfrom x_transformers import TransformerWrapper, Encoder, Decoder\n\nanchor_llm = TransformerWrapper(\n    num_tokens = 20000,\n    max_seq_len = 1024,\n    attn_layers = Decoder(\n        dim = 16,\n        dim_head = 2,\n        depth = 12,\n        heads = 8\n    )\n)\n\naugment_llm1 = TransformerWrapper(\n    num_tokens = 20000,\n    max_seq_len = 1024,\n    attn_layers = Encoder(\n        dim = 16,\n        dim_head = 2,\n        depth = 12,\n        heads = 8\n    )\n)\n\naugment_llm2 = TransformerWrapper(\n    num_tokens = 20000,\n    max_seq_len = 1024,\n    attn_layers = Encoder(\n        dim = 16,\n        dim_head = 2,\n        depth = 12,\n        heads = 8\n    )\n)\n\nvit = ViT(\n    image_size = 256,\n    patch_size = 32,\n    num_classes = 1000,\n    dim = 256,\n    depth = 6,\n    heads = 16,\n    mlp_dim = 2048\n)\n\n# calm\n\nfrom CALM_pytorch import CALM, AugmentParams, FineTuner\n\ncalm = CALM(\n    anchor_llm = anchor_llm,\n    augment_llms = (\n        AugmentParams(\n            model = augment_llm1,\n            mask_kwarg = 'mask'\n        ),\n        AugmentParams(\n            model = augment_llm2,\n            mask_kwarg = 'mask'\n        ),\n        AugmentParams(\n            model = vit,\n            input_shape = (3, 256, 256),\n            hidden_position = 'input',\n            extract_blocks_fn = lambda vit: [m for m in vit.modules() if isinstance(m, Attention)]\n        )\n    ),\n    attn_kwargs = dict(\n        linear_project_context = True,\n        pre_rmsnorm = True,\n        flash = True\n    )\n)\n\nseq = torch.randint(0, 20000, (1, 1024))\nmask = torch.ones((1, 1024)).bool()\n\nprompt = (\n    torch.randint(0, 20000, (1, 256)),\n    torch.randint(0, 20000, (1, 256)),\n    torch.randn(1, 3, 256, 256)\n)\n\nloss = calm(\n    seq,\n    mask = mask,\n    prompt = prompt\n)\n\nloss.backward()\n```\n\n## Todo\n\n- [x] figure out how to correctly mask augment llm tokens\n- [x] auto-derive model dimensions with dummy input\n- [x] take care of finetuning training logic\n- [x] show example of manual definitions of custom connectivity between 2+ attention networks\n- [x] if anchor and augment transformer block modules are directly passed in (without extraction fn), run a dummy input through both networks and order them correctly using hooks\n- [x] fix example for x-transformers, as in x-transformers, depth is actually depth x 2, taking hiddens from after attention and ff\n- [x] when finely specifying hidden positions, make sure to reorder it if the transformer blocks themselves were passed in and not ordered to begin with\n- [x] extend to a list of augmentation llms\n    - [x] full connectivity customization\n    - [x] custom number of augmentation layers per augmetation llm\n    - [x] make simple vit work\n        - [x] refactor so extraction fn, mask kwarg, and other related hparams are grouped together under a dictionary of {[augment_llm_name]: {augment_llm_related_hparams}} - use dataclasses\n        - [x] show example\n- [x] take care of caching the augment hiddens when sampling. forget about anchor kv cache for now\n    - [x] logic for not releasing the saved output from recorder, for inference\n    - [x] managing cross attention block state for popping the saved output from the recorder\n    - [x] move the augmentation forwards into one shared method, and craft out sampling method for anchor\n\n- [ ] able to wire up with just module names\n- [ ] show an example with giving the LLM ability to hear as well, using \u003ca href=\"https://github.com/lucidrains/audiolm-pytorch\"\u003ehubert or wav2vec\u003c/a\u003e wrappers\n- [ ] handle a wrapper or function that takes in the sequence and prompt length, and auto derives the inputs to CALM\n- [ ] add an option for self attention path way with memory tokens attending to hidden states of all augmentation llms, akin to what was done with \u003ca href=\"https://github.com/lucidrains/zorro-pytorch\"\u003eZorro\u003c/a\u003e\n\n## Citations\n\n```bibtex\n@inproceedings{Bansal2024LLMAL,\n  title   = {LLM Augmented LLMs: Expanding Capabilities through Composition},\n  author  = {Rachit Bansal and Bidisha Samanta and Siddharth Dalmia and Nitish Gupta and Shikhar Vashishth and Sriram Ganapathy and Abhishek Bapna and Prateek Jain and Partha Pratim Talukdar},\n  year    = {2024},\n  url     = {https://api.semanticscholar.org/CorpusID:266755751}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fcalm-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fcalm-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fcalm-pytorch/lists"}