burn_pack::Writer requires eager materialization of all tensor bytes before writing

[antimora]

Summary

burn_pack::Writer requires every tensor's Bytes to already be resident in memory before writing starts. write_to_file / into_bytes / size() all call plan() → build_descriptors() first, which reads tensor.bytes.len() for every Tensor in the Vec<Tensor> to compute the aligned offset table — before any I/O happens. So the caller must materialize every tensor's data up front just to build the Vec<Tensor> that goes into Writer::new.

This is asymmetric with the reader: the README advertises tensor data as "read lazily from files," which is true — Reader::from_file reads only the header+metadata up front, and each tensor's bytes are a Bytes::view window read from disk on first access. There's no equivalent for writing.

It's also a regression from the pre-#5064 burn_store::BurnpackWriter, which held Vec<TensorSnapshot> and called snapshot.to_data() inside the per-tensor write loop — materialize, write, drop, repeat — bounding peak memory to the largest single tensor instead of the whole tensor set.

Impact

BurnpackStore::collect_from has the same shape: it maps every TensorSnapshot through bridge::snapshot_to_tensor (which calls to_data(), materializing) before building Vec<Tensor>. Any caller saving a large Module — or, in our case, burn-onnx, which writes .bpk files during codegen straight from a Vec<TensorSnapshot> — must hold the entire tensor set in host memory at once. burn-onnx's model-check suite includes multi-GB models (stable-diffusion-xl, qwen), so this meaningfully raises peak RSS during codegen.

We worked around it in burn-onnx with a hand-rolled streaming writer. Doing that required duplicating burn-pack's private on-disk layout types (Metadata, TensorDescriptor are pub(crate)) — not something a downstream crate should have to maintain.

Root cause

• build_descriptors() reads tensor.bytes.len() for every tensor to compute offsets, before any bytes are written.
• Tensor.bytes: Bytes is a concrete, already-resolved field — no lazy-provider variant.
• write_tensor_data's WRITE_CHUNK_SIZE streaming only bounds one already-resident tensor's write (useful for device-resident tensors); it doesn't address how many tensors are resident simultaneously going into the write loop.

Suggested fix

A two-phase Writer: register each tensor's metadata (name, dtype, shape, byte length, param_id) without its bytes to build the offset table, then supply bytes per tensor via a callback/lazy provider during the actual write pass — mirroring what the old BurnpackWriter did with TensorSnapshot.

Environment

Confirmed against main @ f2321dde9549f5b9b2296b34876e168ebe4efcaf (2026-07-22) — writer.rs, reader.rs, and base.rs are unchanged since #5064 (2026-06-16).

[antimora]

@nathanielsimard @laggui

[antimora]

Traced this down to pinpoint where the eagerness actually lives, since it's not where I first assumed.

Writer's own plumbing is lazy-compatible: build_descriptors() only reads Bytes::len() (a cached field, not a materializing read), and write_tensor_data streams via Bytes::view() one tensor at a time. Fed genuinely lazy Bytes, it wouldn't force everything resident at once.

The actual bottleneck is upstream of Writer, in how bytes get produced in the first place: TensorSnapshot::to_data() → Tensor::to_data() → into_data() → try_read_sync(into_data_async_impl(...)) (burn-tensor/src/tensor/api/base.rs:3310). try_read_sync blocks until the future resolves, and that future is what performs the actual device→host read (ComputeClient::read_async → do_read → server.read() in cubecl-runtime). to_data()'s own doc comment hints at this: "prefer using a Transaction when reading multiple tensors at once — this may improve laziness" — the plain default call (what TensorSnapshot::from_float/from_int/from_bool all use) is fully synchronous and fully materializing; laziness would require deliberately opting into the batched Transaction API, which nothing in the save path does.

So BurnpackStore::collect_from mapping every snapshot through bridge::snapshot_to_tensor (→ to_data()) really does materialize every tensor's bytes into host memory, one after another, before Writer::new is even constructed. There is a genuinely lazy AllocationProperty::Device variant of Bytes in cubecl-runtime ("materializes on access"), but it's used for cross-device staging in ComputeClient::staging, not for what read()/read_async() hand back — those are async specifically because they wait for a real transfer to complete; a lazy handle wouldn't need that.

Net: the format/Writer isn't the bottleneck. It's Tensor::to_data()'s synchronous, fully-materializing default that every snapshot-based save path — BurnpackStore::collect_from and burn-onnx's writer alike — goes through. A fix here would need either a lazy-provider Tensor.bytes on the burn-pack side (as suggested above) and something upstream of it (e.g. TensorSnapshot) that defers the to_data() call until the write loop actually reaches that tensor, rather than resolving it during collection.