metrale_circuit/memory/
weights.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: Weight bytes per node: what the checkpoint stores for it (its declared formats,
4//! scales included) and every load-time copy the loader derives from it ([`super::copies`]).
5//!
6//! Owner: metrale-circuit (memory).
7//! Invariants:
8//! - A weight is `out x k`: `k` the node's first input edge's dim, `out` the sum of its output
9//!   edges' dims. Routed experts hold `experts` such weights, whatever the rows routed to them.
10//! - The embedding table is `vocab x hidden` at its output edge's format, counted once: the draft
11//!   head reads the main table.
12//! - Stored bytes are the declared circuit's (the checkpoint's own formats); served bytes appear
13//!   only through copy rules, since the loader keeps every stored tensor resident.
14//! - A tensor is stored once: a node whose binding another node already holds stores nothing,
15//!   though copies may still derive from it. The draft head's `lm_head` reads the target's table
16//!   (no checkpoint ships its own; Qwen's circuit binds the runtime copy as `mtp.lm_head`), so it
17//!   stores nothing either.
18
19use std::collections::BTreeMap;
20
21use super::MemoryError;
22use super::copies::CopyRule;
23use crate::format::Format;
24use crate::ir::{Circuit, Node, NodeIdx, OpKind, Section};
25
26/// 2026-10-02: One derived copy of a node's weight.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct DerivedCopy {
29    /// 2026-10-02: The rule's id.
30    pub rule: String,
31    /// 2026-10-02: Bytes (all copies, every expert).
32    pub bytes: u64,
33    /// 2026-10-02: Nothing reads it.
34    pub leaked: bool,
35}
36
37/// 2026-10-02: One node's weights.
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct NodeWeights {
40    /// 2026-10-02: The node (in the served circuit).
41    pub node: NodeIdx,
42    /// 2026-10-02: Stored bytes at the declared formats.
43    pub stored: u64,
44    /// 2026-10-02: Derived copies.
45    pub derived: Vec<DerivedCopy>,
46}
47
48/// 2026-10-02: Which section each node of `c` belongs to.
49pub(crate) fn sections(c: &Circuit) -> Vec<Section> {
50    let mut v = vec![Section::Main; c.nodes.len()];
51    for b in &c.blocks {
52        v[b.first..b.end].iter_mut().for_each(|s| *s = b.section);
53    }
54    v
55}
56
57/// 2026-10-02: `(out, k)` of a weight-reading node.
58pub(crate) fn weight_shape(c: &Circuit, n: &Node) -> (u64, u64) {
59    let k = n.inputs.first().map_or(0, |&e| c.edges[e].dim_value);
60    let out = n.outputs.iter().map(|&e| c.edges[e].dim_value).sum();
61    (out, k)
62}
63
64fn copies_per_node(c: &Circuit, n: &Node) -> Result<u64, MemoryError> {
65    match n.op {
66        OpKind::ExpertGateUp | OpKind::ExpertDown => {
67            c.dims
68                .get("experts")
69                .copied()
70                .ok_or_else(|| MemoryError::MissingDim {
71                    node: n.id.clone(),
72                    dim: "experts".into(),
73                })
74        }
75        _ => Ok(1),
76    }
77}
78
79fn bytes_of(n: &Node, what: &str, v: Option<u64>) -> Result<u64, MemoryError> {
80    v.ok_or_else(|| MemoryError::Size {
81        node: n.id.clone(),
82        what: what.to_string(),
83    })
84}
85
86/// 2026-10-02: Every weight of `served` (its embedding and every weight-reading node), stored at
87/// the formats `declared` gives the node of the same id, with the copies `rules` derive under
88/// `settings`.
89pub fn node_weights(
90    served: &Circuit,
91    declared: &Circuit,
92    rules: &[CopyRule],
93    settings: &BTreeMap<String, String>,
94) -> Result<Vec<NodeWeights>, MemoryError> {
95    let secs = sections(served);
96    let declared_of: BTreeMap<&str, &Node> =
97        declared.nodes.iter().map(|n| (n.id.as_str(), n)).collect();
98    let mut out = Vec::new();
99    let mut bound: std::collections::BTreeSet<&[String]> = std::collections::BTreeSet::new();
100    for (i, n) in served.nodes.iter().enumerate() {
101        if n.op == OpKind::Embed {
102            if secs[i] != Section::Main {
103                continue;
104            }
105            let vocab = *served
106                .dims
107                .get("vocab")
108                .ok_or_else(|| MemoryError::MissingDim {
109                    node: n.id.clone(),
110                    dim: "vocab".into(),
111                })?;
112            let e = n.outputs.first().map(|&e| &served.edges[e]);
113            let table = e.and_then(|e| e.format.bytes(vocab, e.dim_value));
114            out.push(NodeWeights {
115                node: i,
116                stored: bytes_of(n, "embedding table", table)?,
117                derived: Vec::new(),
118            });
119            continue;
120        }
121        let Some(serve_fmt) = n.weight else { continue };
122        let stored_fmt: Format = declared_of
123            .get(n.id.as_str())
124            .and_then(|d| d.weight)
125            .unwrap_or(serve_fmt);
126        let (o, k) = weight_shape(served, n);
127        let copies = copies_per_node(served, n)?;
128        let one = bytes_of(n, "stored weight", stored_fmt.weight_bytes(o, k))?;
129        let draft_head = secs[i] == Section::Draft && n.op == OpKind::LmHead;
130        let first = !draft_head && (n.binding.is_empty() || bound.insert(n.binding.as_slice()));
131        let stored = match first {
132            true => bytes_of(n, "stored weight", one.checked_mul(copies))?,
133            false => 0,
134        };
135        let mut derived = Vec::new();
136        for r in rules.iter().filter(|r| {
137            r.matches(
138                &served.arch,
139                &n.op,
140                secs[i],
141                stored_fmt,
142                serve_fmt,
143                settings,
144            )
145        }) {
146            let b = r
147                .law
148                .bytes(o, k)
149                .and_then(|b| b.checked_mul(copies))
150                .and_then(|b| b.checked_mul(r.count));
151            derived.push(DerivedCopy {
152                rule: r.id.clone(),
153                bytes: bytes_of(n, &format!("copy `{}`", r.id), b)?,
154                leaked: r.leaked,
155            });
156        }
157        out.push(NodeWeights {
158            node: i,
159            stored,
160            derived,
161        });
162    }
163    Ok(out)
164}
metrale_circuit/memory/weights.rs

metrale_circuit/memory/
weights.rs
