huggingface / huggingface/candle

`candle-transformers::models::llama`: expose `Block` for per-block composition

Open
#3,817 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Rust
Stars
21.1k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

Problem

Llama::load_with_mlp_factory (#3809) closed the case where a caller wants to
change a block's MLP and inherit everything else. It does not help a caller
that needs less than a whole model.

Pipeline parallelism is the case. A stage is a contiguous range of transformer
layers placed on one device: stage 0 owns the token embedding, the last stage
owns the final norm and the LM head, and every stage in between owns neither
— just its blocks, its slice of the KV cache, and the activations handed to it
by the previous stage. There is no Llama to construct: load_with_mlp_factory
builds embeddings, all num_hidden_layers blocks, the norm and the head, which
is precisely what a stage must not materialise. A four-stage deployment that
called it would load the full model four times.

The same gap appears one step away, in mixture-of-experts. A checkpoint whose
layers are individually dense or MoE needs attention composed with a per-layer
choice of feed-forward. The MLP factory's usize argument expresses the choice
— but only inside a whole-model load, so a caller that wants the block itself
still has to build one.

Both therefore reimplement the block: RmsNorm → attention → RmsNorm → MLP,
plus the causal mask, the rotary tables and the KV cache that attention needs.
In Tachyon-Mesh that is ~269 lines of copied attention/cache/mask math, kept
honest only by a test asserting bit-for-bit equality against the dense
reference. The MLP seam did not reduce it, because what these callers duplicate
is not the MLP — it is everything around it.

What makes this worth a second look is how close the code already is.
Block::load_with_mlp is what load_with_mlp_factory calls per layer, and
Cache — rotary tables, masks, per-layer KV — is already public and already
indexed by block_idx. The composition these callers need is the one the
crate performs internally; only the visibility differs.

What a usable seam needs

  1. Construct one block, at a known layer index, from that layer's
    VarBuilder — with the same MLP substitution load_with_mlp_factory
    already offers, so the two seams compose rather than compete.
  2. Run it against the public Cache. Block::forward already takes
    block_idx and indexes the cache with it, which is exactly what a stage
    holding layers 12..18 needs; no new cache shape is required.
  3. Require neither embedding nor head. A middle stage has no wte and no
    lm_head, and must not be made to load them.
  4. Stay additive. Llama::load, load_with_mlp_factory and forward must
    behave exactly as they do now — the shape #3721 and #3809 established.

Proposed API

Visibility, not new machinery:

pub struct Block { /* … */ }

impl Block {
    /// Load layer `layer_idx`, with the dense MLP.
    pub fn load(vb: VarBuilder, cfg: &Config, layer_idx: usize) -> Result<Self>;

    /// Load layer `layer_idx`, substituting the feed-forward computation —
    /// the per-layer half of `Llama::load_with_mlp_factory`.
    pub fn load_with_block_mlp(
        vb: VarBuilder,
        cfg: &Config,
        layer_idx: usize,
        mlp: Box<dyn BlockMlp>,
    ) -> Result<Self>;

    pub fn forward(
        &self,
        x: &Tensor,
        index_pos: usize,
        block_idx: usize,
        cache: &mut Cache,
        adapters: Option<&[Option<&str>]>,
    ) -> Result<Tensor>;
}

load_with_block_mlp takes the trait object rather than the private
BlockMlpImpl, so the enum stays an implementation detail; load_with_mlp_factory
can keep using the enum internally for the dense case.

Note

This is the second half of #3809. That one let a caller replace a block's MLP
inside a whole model; this one lets a caller assemble blocks without one. With
both, a tensor-parallel model, a pipeline stage and a mixed dense/MoE stack are
all compositions of upstream parts, and the copied attention goes away.

I am happy to submit the PR, with the equality test against the dense reference
extended to a block assembled this way. Related: #3809, #3721.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with Block::load_with_mlp, Block::forward, Cache, and Llama::load_with_mlp_factory to understand the existing per-layer construction and cache indexing. Extend the equality test against the dense reference, and consider the work complete when standalone blocks can run with the public Cache while existing Llama loading and forward behavior remain unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
58/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.