PretrainedReader.ixx File Reference

Reader for Mila pretrained binary format. More...

#include <cstdio>
#include <filesystem>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <format>
#include <algorithm>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <exception>
#include <type_traits>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
import Compute.Device;
import Dnn.TensorBuffer;
import Dnn.Tensor;
import Dnn.TensorDataTypeTraits;
import Serialization.SafeTensors;
import Compute.CpuMemoryResource;
import Serialization.Tensor;
import Dnn.TensorTypes;
import Dnn.ITensor;
import Serialization.OpenMode;
import Serialization.Serializer;
import Dnn.TensorDataType;
import nlohmann.json;

Classes
struct	Mila::Dnn::Serialization::PretrainedMetadata
	Metadata for pretrained model. More...
class	Mila::Dnn::Serialization::PretrainedModelReader
	Reader for Mila pretrained binary format. More...
struct	Mila::Dnn::Serialization::TensorBlobMetadata
	Metadata for a tensor blob in pretrained model format. More...

Namespaces
namespace	Mila
	Mila main API namespace.

Enumerations
enum class	DType : uint32_t {   Float32 = 0 , Float16 = 1 , BFloat16 = 2 , Int32 = 3 ,   UInt8 = 4 , Float8E4M3 = 5 , Float8E5M2 = 6 , Int8 = 7 }

Functions
TensorDataType	Mila::Dnn::Serialization::dtypeToTensorDataType (uint32_t dtype)
uint32_t	Mila::Dnn::Serialization::tensorDataTypeToWireCode (TensorDataType type)
std::string	Mila::Dnn::Serialization::toMetadataJSON (const PretrainedMetadata &metadata)
	Serialize PretrainedMetadata to the JSON the reader parses back.

Detailed Description

Reader for Mila pretrained binary format.

Provides direct access to pretrained model weights stored in Mila's flat binary format. Used by fromPretrained() factory methods.

The whole file is memory-mapped at construction (CreateFileMapping/MapViewOfFile on Windows, mmap on POSIX). streamTensorBlobs() consumes blobs in ascending file offset order, turning the former 224+ random per-tensor reads into a single sequential scan. On the CUDA path a background producer thread stages each blob into a pinned host buffer (double-buffered) so disk I/O overlaps the H2D copy; all CUDA work stays on the consuming thread. Staging reads each blob directly from the file handle (positioned ReadFile/pread), not by faulting through the mapped view – the map's 4 KB on-demand faults throttle a large model well below disk bandwidth once the file no longer fits resident alongside the model's own GPU staging. The mapping is kept for oversized blobs (consumed straight from the view) and the legacy random-access readTensorBlob<MR>() fallback.

Function Documentation

toMetadataJSON()

std::string Mila::Dnn::Serialization::toMetadataJSON ( const PretrainedMetadata & metadata )

inlineexport

Serialize PretrainedMetadata to the JSON the reader parses back.

The inverse of parseMetadataJSON, which until now had none – an artifact could be inspectable without being loadable. Every field the parser extracts is emitted, so a written model carries the same architecture description a converted one does.

Keys are quoted on both sides by the parser, so no key can match inside a longer one ("rope_theta" does not match within "rope_theta_local"). Do not introduce a key that is a prefix of another up to its closing quote.