Mila::Dnn::Serialization::PretrainedModelReader Class Reference

Reader for Mila pretrained binary format. More...

Public Member Functions
	PretrainedModelReader (const std::filesystem::path &filepath)
	Open a Mila model file for reading.
bool	close ()
const std::string &	getFilename () const noexcept
size_t	getMaxTensorSizeBytes () const
	Get the maximum byte size across all tensors in the index.
const PretrainedMetadata &	getPretrainedMetadata () const
	Get pretrained model metadata.
std::vector< std::string >	getTensorNames () const
	Get list of all tensor names in the model.
size_t	getTensorSizeBytes (const std::string &name) const
	Get the raw byte size of a named tensor.
const std::string &	getWeightQuantization () const noexcept
	Weight quantization the artifact was written with, empty if unquantized.
bool	hasTensor (const std::string &name) const
	Check if tensor exists.
bool	isOpen () const noexcept
template<typename MR = Compute::CpuMemoryResource> requires isValidTensor<dtype_t::UINT8, MR>
TensorBlob< MR >	readTensorBlob (const std::string &name, int device_id=0)
	Read raw tensor bytes by name into a memory-resource-typed blob.
template<typename TStagingMemoryResource = Compute::CpuMemoryResource, typename TConsumer>
void	streamTensorBlobs (TConsumer &&consume, int device_id=0)
	Stream every tensor blob to a consumer in file-offset order.

Detailed Description

Reader for Mila pretrained binary format.

Two containers are accepted, sniffed by the leading magic. Both fill the same tensor index, so everything past the header parse – the mapping, the offset-ordered stream, the pinned staging producer – is common.

MILA (every .bin already on disk; support for it is permanent):

• Header: MILA magic (0x4D494C41), version, num_tensors
• Metadata: JSON string with model configuration
• Tensor index: for each tensor: name, dtype, shape, offset, nbytes
• Tensor data: concatenated binary blobs

safetensors (what Mila now writes):

• Header: u64 little-endian header length, then that many bytes of JSON
• Each JSON entry: dtype, shape, data_offsets relative to the data region
• Model configuration rides in metadata under the mila_config key

Provides flat key-value access to tensors by name:

• "lenc.wte.weight"
• "tf_layer_0.ln_1.bias"
• "ln_final.weight"

Usage:

PretrainedModelReader reader( "gpt2_small.bin" );
 
auto metadata = reader.getPretrainedMetadata();
auto names = reader.getTensorNames();
 
for (const auto& name : names)
{
    auto blob = reader.readTensorBlob<CudaPinnedMemoryResource>( name, device_id );
    network->loadTensorByFlatName( name, blob );
}

Constructor & Destructor Documentation

PretrainedModelReader()

Mila::Dnn::Serialization::PretrainedModelReader::PretrainedModelReader ( const std::filesystem::path & filepath )

inlineexplicit

Open a Mila model file for reading.

Parameters

filepath

Path to .bin model file.

Exceptions

std::runtime_error

if file cannot be opened or format is invalid.

Member Function Documentation

getMaxTensorSizeBytes()

size_t Mila::Dnn::Serialization::PretrainedModelReader::getMaxTensorSizeBytes ( ) const

inline

Get the maximum byte size across all tensors in the index.

Returns the largest nbytes value in the tensor index. All sizes are known at construction time. No I/O is performed.

Returns

size_t Maximum tensor byte count, or 0 if the index is empty.

getTensorSizeBytes()

size_t Mila::Dnn::Serialization::PretrainedModelReader::getTensorSizeBytes ( const std::string & name ) const

inline

Get the raw byte size of a named tensor.

All sizes are known at construction time from the tensor index. No I/O is performed.

Parameters

name	Tensor name.

Returns

size_t Byte count of the tensor data.

Exceptions

std::runtime_error

if name is not found.

getWeightQuantization()

const std::string & Mila::Dnn::Serialization::PretrainedModelReader::getWeightQuantization ( ) const

inlinenoexcept

Weight quantization the artifact was written with, empty if unquantized.

Only a pre-quantized artifact carries this. A MILA .bin and a BF16 safetensors file both return empty, which means "quantize on load" – the behaviour that predates pre-quantized artifacts.

The value distinguishes policies a dtype cannot: FP4 at group 128 and group 64 are both packed into U8, so only this string can refuse the wrong one.

readTensorBlob()

template<typename MR = Compute::CpuMemoryResource>
requires isValidTensor<dtype_t::UINT8, MR>

TensorBlob< MR > Mila::Dnn::Serialization::PretrainedModelReader::readTensorBlob ( const std::string & name, int device_id = 0 )

inline

Read raw tensor bytes by name into a memory-resource-typed blob.

Allocates a TensorBuffer<UINT8, MR> of the exact tensor byte size and reads directly from the file into it. No intermediate buffer is used. When MR is CudaPinnedMemoryResource the returned blob data is page-locked, enabling direct DMA to device in copyFromBlob without a staging copy.

Template Parameters

MR	Memory resource for the blob data buffer. Defaults to CpuMemoryResource.

Parameters

name	Tensor name.
device_id	Device index passed to the memory resource constructor.

Returns

TensorBlob<MR> owning the metadata and raw byte buffer.

Exceptions

std::runtime_error

if the tensor is not found or the read fails.

streamTensorBlobs()

template<typename TStagingMemoryResource = Compute::CpuMemoryResource, typename TConsumer>

void Mila::Dnn::Serialization::PretrainedModelReader::streamTensorBlobs ( TConsumer && consume, int device_id = 0 )

inline

Stream every tensor blob to a consumer in file-offset order.

Replaces the per-tensor seek+read loop. Because the whole file is mapped once, consuming in ascending offset is a single sequential scan the OS can read ahead, rather than 224+ random reads in hash-map order.

When TStagingMemoryResource is CudaPinnedMemoryResource a background producer thread stages each blob mmap -> pinned host buffer (double-buffered) while the calling thread runs the consumer (H2D + quantize). All CUDA calls stay on the calling thread; the producer does only host memcpy, matching the safe split in TokenSequenceLoader. Blobs larger than the staging buffer (e.g. the token embedding) bypass staging and are consumed directly from the mapped view.

Contract: consume() MUST complete every device read of blob.data() before it returns, because the pinned slot is handed back to the producer for reuse the moment consume() returns. The non-quantized copyFromBlob path self-synchronizes on the default stream, but the FP8/FP4 quantize path issues an async H2D on the op stream and does NOT, so the model's consume callback must synchronize its execution context after loadParameter. The producer's next memcpy overlaps that synchronize, preserving the disk/H2D overlap.

Template Parameters

TStagingMemoryResource

Staging resource. CudaPinnedMemoryResource selects the threaded pinned path; CpuMemoryResource consumes mapped views directly with no staging and no producer thread.

Parameters

consume	Callable invoked as consume(const std::string& name, const ITensorBlob&).
device_id	Device index for the pinned staging buffers (CUDA path only).

---

The documentation for this class was generated from the following file:

• Mila/Src/Dnn/Serialization/PretrainedReader.ixx