Skip to content

Incremental Adapter Composition (Add / Remove / Update) #21

@antonpibm

Description

@antonpibm

Feature Request: Incremental Adapter Composition (Add / Remove / Update)

Problem

Currently GraniteSwitchComposer.from_base_and_adapters() only supports building a fresh Granite Switch model from a base model + a full set of adapters. There is no way to take an already-composed model and add, remove, or update individual adapters without repeating the entire composition from scratch.

This makes iteration cumbersome — updateing an adapter in existing model requires a full rebuild, this might not even be possible if the original composition setup is not available.

Proposed Capability

Support three incremental operations on an existing composed model:

Operation Input Effect
Add Existing composed model + new adapter(s) Appends adapter slot(s), expands stacked tensors, adds control tokens
Remove Existing composed model + adapter name(s) Removes adapter slot(s), compacts tensors, removes control tokens
Update Existing composed model + updated adapter(s) Overwrites adapter weights in-place, no structural change

Current Architecture Constraints

The compose pipeline makes several assumptions that tie it to fresh builds:

  1. Fixed-size stacked tensors — LoRA weights are stored as [num_adapters, 1, ...] tensors. Adding/removing adapters requires reshaping every LoRA parameter in the model.
  2. Positional adapter indices — Adapter i always occupies slot i in stacked tensors, adapter_token_ids[i], adapter_names[i], etc. Removing an adapter from the middle creates index gaps.
  3. Config is staticGraniteSwitchConfig is built once; num_adapters, adapter_token_ids, hiding_groups, and hiding_policy are all computed at compose time.
  4. Tokenizer is additive — Control tokens are added but never removed. The tokenizer vocabulary only grows.
  5. Base weights are re-transferredtransfer_base_weights loads the upstream base model every time, even though base weights are identical across compositions.

High-Level Implementation Approach

Phase 1: Load Existing Composed Model

Add a class method to load an already-composed checkpoint:

GraniteSwitchComposer.from_existing(model_path) -> (model, config, tokenizer, adapter_registry)

This loads the model via GraniteSwitchForCausalLM.from_pretrained() and reconstructs the adapter registry (name → slot index mapping) from the config.

Phase 3: Implement Operations

Add Adapter

  1. Load existing model + registry
  2. Discover and validate new adapter(s) — rank must be ≤ max_lora_rank (or resize all if larger)
  3. Expand stacked LoRA tensors along dim 0: [N, 1, ...] → [N+K, 1, ...]
  4. Remap and insert new adapter weights into the new slots
  5. Add control tokens to tokenizer, resize embeddings
  6. Update config: num_adapters, adapter_token_ids, adapter_names, hiding_groups, hiding_policy
  7. Re-validate and save

Remove Adapter

  1. Load existing model + registry
  2. Resolve adapter name(s) to slot indices
  3. Delete rows from stacked LoRA tensors (compact remaining)
  4. Remap indices in remaining config fields
  5. Remove control tokens from tokenizer (or mark as unused — see open question)
  6. Update config
  7. Re-validate and save

Update Adapter

  1. Load existing model + registry
  2. Resolve adapter name to slot index
  3. Load new adapter weights, validate rank/modules match existing slot
  4. Zero the slot in all stacked tensors, then write new weights in-place
  5. Update registry metadata (source path, timestamp)
  6. Re-validate and save

Phase 4: CLI Integration Proposition

Extend the compose_granite_switch CLI:

# Fresh build (existing behavior, unchanged)
python -m granite_switch.composer.compose_granite_switch \
  --adapters adapter1 adapter2

# Add adapter(s) to existing model
python -m granite_switch.composer.compose_granite_switch \
  --model ./existing-composed-model \
  --add-adapters new_adapter1 new_adapter2

# Remove adapter(s)
python -m granite_switch.composer.compose_granite_switch \
  --model ./existing-composed-model \
  --remove-adapters adapter_name

# Update adapter(s) in-place
python -m granite_switch.composer.compose_granite_switch \
  --model ./existing-composed-model \
  --update-adapters adapter_name=/path/to/new/weights

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions