huggingface / huggingface/candle
`candle-transformers::models::llama`: expose `Block` for per-block composition
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
- Construct one block, at a known layer index, from that layer's
VarBuilder— with the same MLP substitutionload_with_mlp_factory
already offers, so the two seams compose rather than compete. - Run it against the public
Cache.Block::forwardalready takes
block_idxand indexes the cache with it, which is exactly what a stage
holding layers 12..18 needs; no new cache shape is required. - Require neither embedding nor head. A middle stage has no
wteand no
lm_head, and must not be made to load them. - Stay additive.
Llama::load,load_with_mlp_factoryandforwardmust
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
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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