Mila
Deep Neural Network Library
Loading...
Searching...
No Matches
Mila::Dnn::GptModel< TDeviceType, TPrecision > Class Template Referenceexport

GPT inference model. More...

Inheritance diagram for Mila::Dnn::GptModel< TDeviceType, TPrecision >:
Mila::Dnn::LanguageModel< TDeviceType, TPrecision > Mila::Dnn::Model< TDeviceType, TPrecision >

Public Types

using GptTransformerType = GptTransformer<TDeviceType, TPrecision>
using ModelBase = LanguageModel<TDeviceType, TPrecision>
using MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using TensorType = Tensor<TPrecision, MR>
using TokenIndexType = Tensor<dtype_t::INT32, MR>
Public Types inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
using Base = Model<TDeviceType, TPrecision>
Public Types inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
using NetworkType = Network<TDeviceType, TPrecision>

Public Member Functions

 GptModel (const GptModel &)=delete
 GptModel (GptModel &&)=default
const GptConfig & getConfig () const noexcept
GptModel & operator= (const GptModel &)=delete
GptModel & operator= (GptModel &&)=default
void saveCheckpoint (const std::filesystem::path &path, SerializationMode mode=SerializationMode::Checkpoint) const
 Write a Mila-native archive that fromCheckpoint() can restore.
std::string toString () const override
 Human-readable summary of this model's configuration.
Public Member Functions inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
 LanguageModel (const LanguageModel &)=delete
 LanguageModel (LanguageModel &&)=default
std::vector< std::string > componentPaths () const
 Every observable component path, for choosing a pattern or diagnosing a zero.
GenerateStatus generate (std::span< const int32_t > prompt_tokens, const std::function< void(int32_t)> &on_token, const GenerateParams &params={}, std::stop_token stop={})
 Generate tokens from a prompt, streaming each through on_token.
size_t observe (std::string_view pattern, Compute::ComputePassMask passes, Compute::ActivationObserver sink)
 Watch activations flowing through matching components on matching passes.
LanguageModel & operator= (const LanguageModel &)=delete
LanguageModel & operator= (LanguageModel &&)=default
void save (const std::filesystem::path &path) const
 Write this model's live weights as a safetensors weights file.
void seedSampler (uint64_t seed)
 Seed the sampler's RNG for reproducible generation.
void stopObserving ()
 Detach every observer and clear the sink. A probe that outlives its question costs.
std::string weightQuantizationScheme () const
 The scheme name save records for these weights, e.g.
Public Member Functions inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
 Model (const Model &)=delete
 Model (Model &&)=default
DeviceId getDeviceId () const noexcept
 The device this model runs on.
MemoryStats getMemoryStats () const
 Current memory allocation breakdown for this model.
RuntimeMode getRuntimeMode () const noexcept
 The runtime mode this model was constructed for.
bool isEval () const noexcept
 True if this model is currently in eval sub-state.
bool isInferenceMode () const noexcept
 True if this model was constructed for inference.
bool isTrainingMode () const noexcept
 True if this model was constructed for training.
Model & operator= (const Model &)=delete
Model & operator= (Model &&)=default
void setEval (bool eval)
 Toggle eval sub-state for this model.
void train ()
 Run the training loop for this model.

Static Public Member Functions

static std::unique_ptr< GptModel > fromCheckpoint (const std::filesystem::path &path, DeviceId device_id=DeviceId{ TDeviceType, 0 }, dim_t context_length=0)
 Load from a Mila-native serialized artifact.
static std::unique_ptr< GptModel > load (const std::filesystem::path &path, dim_t context_length, DeviceId device_id=DeviceId{ TDeviceType, 0 }, bool strict=true)
 Load GPT-2 from a Mila weights file.

Protected Member Functions

int32_t eosToken () const noexcept override
 GPT-2 end-of-text token id.
dim_t maxSequenceLength () const noexcept override
 Maximum sequence length from GPT config.
GenerateStatus onGenerating (std::span< const int32_t > prompt_tokens, const std::function< void(int32_t)> &on_token, const GenerateParams &params, std::stop_token stop) override
 Prefill + KV-cache decode loop with per-token streaming.
void onTraining () override
 Training loop – not yet implemented for GptModel.
dim_t vocabSize () const noexcept override
 Vocabulary size from GPT config.
Protected Member Functions inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
 LanguageModel (std::unique_ptr< LanguageModelNetwork< TDeviceType, TPrecision > > network, RuntimeMode runtime_mode, Serialization::WeightsMetadata source_metadata={}, WeightQuantization weight_quantization=WeightQuantization::None, int fp4_group_size=128)
int32_t awaitSampledToken ()
 Block until the last enqueueSampleNext()'s token id is host-visible.
void enqueueSampleNext (const TensorType &logits, TokenTensor &token_out, const SamplingParams &params)
 Enqueue a sampling step without waiting for the host readback.
virtual float finalLogitSoftcap () const noexcept
 Optional final-logit softcap the sampler applies (0 disables).
const LanguageModelNetwork< TDeviceType, TPrecision > & getNetwork () const noexcept
LanguageModelNetwork< TDeviceType, TPrecision > & getNetwork () noexcept
int32_t sampleNext (const TensorType &logits, TokenTensor &token_out, const SamplingParams &params)
 Sample the next token from a logits row on the device.
virtual std::unordered_set< int32_t > stopTokens () const
Protected Member Functions inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
 Model (std::unique_ptr< NetworkType > network, RuntimeMode runtime_mode)
 Construct with a fully built network and runtime mode.

Additional Inherited Members

Protected Types inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
using MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using TensorType = Tensor<TPrecision, MR>
using TokenTensor = Tensor<TensorDataType::INT32, MR>
Protected Attributes inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
int fp4_group_size_ { 128 }
 The FP4 group the live weights were built at; meaningful only when they are FP4.
Serialization::WeightsMetadata source_metadata_
 The loaded weights file's metadata, carried so save can write it back verbatim.
WeightQuantization weight_quantization_ { WeightQuantization::None }
 What the live weights are, not what the source file was.
Protected Attributes inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
std::unique_ptr< NetworkType > network_
 The owned Network instance.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TPrecision>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::GptModel< TDeviceType, TPrecision >

GPT inference model.

Owns a loaded, built GptTransformer and exposes generate() for autoregressive text generation.

Construction is only possible via load() or fromCheckpoint(). The network is always in a built, weights-loaded, inference-mode state when generation is called.

Member Function Documentation

◆ eosToken()

template<DeviceType TDeviceType, TensorDataType TPrecision>
int32_t Mila::Dnn::GptModel< TDeviceType, TPrecision >::eosToken ( ) const
inlineoverrideprotectedvirtualnoexcept

GPT-2 end-of-text token id.

Implements Mila::Dnn::LanguageModel< TDeviceType, TPrecision >.

◆ fromCheckpoint()

template<DeviceType TDeviceType, TensorDataType TPrecision>
std::unique_ptr< GptModel > Mila::Dnn::GptModel< TDeviceType, TPrecision >::fromCheckpoint ( const std::filesystem::path & path,
DeviceId device_id = DeviceId{ TDeviceType, 0 },
dim_t context_length = 0 )
inlinestatic

Load from a Mila-native serialized artifact.

Reads a checkpoint or weights-only artifact produced by GptTransformer::save() via ModelArchive.

Parameters
pathPath to the Mila archive.
device_idTarget device.
context_lengthDeployment context length. Zero takes the geometry the checkpoint was built with, falling back to the trained maximum.
Returns
Inference-ready GptModel.

◆ load()

template<DeviceType TDeviceType, TensorDataType TPrecision>
std::unique_ptr< GptModel > Mila::Dnn::GptModel< TDeviceType, TPrecision >::load ( const std::filesystem::path & path,
dim_t context_length,
DeviceId device_id = DeviceId{ TDeviceType, 0 },
bool strict = true )
inlinestatic

Load GPT-2 from a Mila weights file.

Reads a Mila weights file produced by converting third-party checkpoints (e.g. HuggingFace GPT-2) via WeightsReader.

Parameters
pathPath to the weights file.
context_lengthMaximum sequence length to build for.
device_idTarget device.
strictThrows on unknown parameter names if true.
Returns
Inference-ready GptModel.

◆ maxSequenceLength()

template<DeviceType TDeviceType, TensorDataType TPrecision>
dim_t Mila::Dnn::GptModel< TDeviceType, TPrecision >::maxSequenceLength ( ) const
inlineoverrideprotectedvirtualnoexcept

Maximum sequence length from GPT config.

Implements Mila::Dnn::LanguageModel< TDeviceType, TPrecision >.

◆ onGenerating()

template<DeviceType TDeviceType, TensorDataType TPrecision>
GenerateStatus Mila::Dnn::GptModel< TDeviceType, TPrecision >::onGenerating ( std::span< const int32_t > prompt_tokens,
const std::function< void(int32_t)> & on_token,
const GenerateParams & params,
std::stop_token stop )
inlineoverrideprotectedvirtual

Prefill + KV-cache decode loop with per-token streaming.

Implements Mila::Dnn::LanguageModel< TDeviceType, TPrecision >.

◆ onTraining()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::GptModel< TDeviceType, TPrecision >::onTraining ( )
inlineoverrideprotectedvirtual

Training loop – not yet implemented for GptModel.

Exceptions
std::runtime_erroralways.

Implements Mila::Dnn::Model< TDeviceType, TPrecision >.

◆ saveCheckpoint()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::GptModel< TDeviceType, TPrecision >::saveCheckpoint ( const std::filesystem::path & path,
SerializationMode mode = SerializationMode::Checkpoint ) const
inline

Write a Mila-native archive that fromCheckpoint() can restore.

Writes the network config, the component graph, and one blob per parameter. Weights only – optimizer state belongs to the trainer, which owns its own archive scope.

Parameters
pathDestination archive path (overwritten if it exists).
modeSerialization mode recorded in the archive.
Exceptions
std::runtime_errorif the archive cannot be opened or a component cannot serialize its parameters.

◆ toString()

template<DeviceType TDeviceType, TensorDataType TPrecision>
std::string Mila::Dnn::GptModel< TDeviceType, TPrecision >::toString ( ) const
inlineoverridevirtual

Human-readable summary of this model's configuration.

Implements Mila::Dnn::Model< TDeviceType, TPrecision >.

◆ vocabSize()

template<DeviceType TDeviceType, TensorDataType TPrecision>
dim_t Mila::Dnn::GptModel< TDeviceType, TPrecision >::vocabSize ( ) const
inlineoverrideprotectedvirtualnoexcept

Vocabulary size from GPT config.

Implements Mila::Dnn::LanguageModel< TDeviceType, TPrecision >.


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