Pure token embedding component (device-templated). More...

Inheritance diagram for Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >:

Public Types
using	ComponentBase = Component<TDeviceType, TPrecision>
using	EmbeddingTensorType = Tensor<TPrecision, MR>
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TableScaleTensorType = Tensor<TTableQuantization::kScaleDtype, MR>
using	TableTensorType = Tensor<kTableDtype, MR>
using	TokenIndexType = Tensor<TIndex, MR>

Public Member Functions
	TokenEmbedding (const std::string &name, const TokenEmbeddingConfig &config, std::optional< DeviceId > device_id=std::nullopt)
	Construct a TokenEmbedding component.
TokenIndexType &	backward (const TokenIndexType &input, const EmbeddingTensorType &output_grad)
	Backward pass – accumulates gradients into wte.
EmbeddingTensorType &	forward (const TokenIndexType &input)
	Forward pass – returns component-owned embeddings tensor.
DeviceId	getDeviceId () const override
	Get the compute device id associated with this component.
int64_t	getEmbeddingDim () const noexcept
std::vector< ITensor * >	getGradients () const override
	Return non-owning pointers to parameter gradient tensors.
MemoryStats	getMemoryStats () const override
	Return the current memory allocation breakdown for this component.
std::vector< std::string >	getParameterNames () const override
	List all available parameter names for this component.
std::vector< ITensor * >	getParameters () const override
	Return non-owning pointers to parameter tensors.
MemoryStats	getRequiredMemory (const BuildContext &context) const override
	What onBuilding() would allocate for this context, without allocating.
const ComponentType	getType () const override
	Get the component type identifier.
int64_t	getVocabSize () const noexcept
std::shared_ptr< TableScaleTensorType >	getWeightScalesTensorShared () const noexcept
std::shared_ptr< TableTensorType >	getWeightTensorShared () const noexcept
EmbeddingTensorType *	getWteGrad () const noexcept
void	loadParameter (const std::string &name, const ITensorBlob &blob) override
	Load a parameter from serialized tensor data.
dim_t	parameterCount () const override
	Return number of trainable parameters.
void	save_ (ModelArchive &archive, SerializationMode mode) const override
void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const override
	Emit the table, and its per-row scales when the table is quantized.
void	synchronize () override
std::string	toString () const override
	Produce a short, human-readable description of the component.
void	zeroGradients () override
	Clear all model-owned gradients for this component.
Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >
	Component (const std::string &name)
	Construct component with required name identifier.
virtual void	build (const BuildContext &context) final
	Build the component with the provided BuildContext (canonical overload).
const std::string	getName () const
	Get the component's name identifier.
TrainingMode	getTrainingMode () const noexcept
	The current runtime behavioral mode of this Component.
virtual bool	isBuilt () const final
	Returns true if build() has completed successfully.
virtual void	load_ (ModelArchive &archive, SerializationMode mode)
	Restore this component's parameters from its archive scope.
virtual void	requireSerializableParameters () const
	Verify this component can serialize whatever parameters it owns.
void	setTrainingMode (TrainingMode mode)
	Set the runtime behavioral mode for this Component.

Static Public Attributes
static constexpr bool	kIsQuantized = TTableQuantization::kIsQuantized
static constexpr TensorDataType	kTableDtype

Protected Member Functions
void	onBuilding (const BuildContext &build_context) override
	Hook invoked by build() to allocate component buffers.
void	onExecutionContextSet () override
	Lifecycle hook: Called immediately after ExecutionContext is set.
void	onTrainingModeChanging (TrainingMode training_mode) override
	Hook called before TrainingMode transitions.
Protected Member Functions inherited from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >
IExecutionContext *	getExecutionContext () const
	Get the shared execution context.
bool	hasExecutionContext () const noexcept
	Check if execution context has been set.
void	loadParameterFromBlob (const std::string &param_name, const Serialization::ITensorBlob &blob, Tensor< TParameterPrecision, TMemoryResource > &target, const shape_t &expected_shape)
	Load a tensor blob into a parameter tensor with validation.
void	saveParameterToArchive (ModelArchive &archive, const std::string &parameter_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter) const
	Write one parameter tensor into the archive under "tensors/<name>".
void	saveParameterToWriter (Serialization::SafeTensorsWriter &writer, const std::string &flat_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter, Serialization::TensorSavePass pass) const
	Drive one parameter through one pass of a flat safetensors save.
void	setExecutionContext (IExecutionContext *context)
	Set the execution context for this component.

Additional Inherited Members
Static Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >
static constexpr DeviceType	getDeviceType ()
	Compile-time device type for this component instance.
static constexpr TensorDataType	getPrecision () noexcept
	Compile-time tensor precision for this component instance.
Protected Types inherited from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >
using	HostStagingMemoryResource
	Host memory a device-resident parameter stages through.
Protected Attributes inherited from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >
BuildContext	build_context_
	The BuildContext stored at build time.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >

Pure token embedding component (device-templated).

Transforms input token indices into continuous vector representations by looking up each index in the vocabulary embedding table (wte). No positional information is added here.

TTableQuantization = PerChannelFp8<> stores the table as FP8_E4M3 with one float32 absmax scale per vocabulary row, quantized at loadParameter() time (D4 Design B). The vocab-row scale axis coincides with a tied lm_head's per-output-channel scale axis, so getWeightTensorShared() plus getWeightScalesTensorShared() feed Linear::installSharedWeight directly. The quantized path is inference-only.

Construction modes:

Standalone: provide a DeviceId to create and own an ExecutionContext.
Child/deferred: omit DeviceId; caller must call setExecutionContext() before build().

Template Parameters

TDeviceType	Device type (DeviceType::Cpu or DeviceType::Cuda).
TIndex	Data type for token indices (typically INT32).
TPrecision	Tensor precision for embeddings (FP32 or BF16).
TTableQuantization	Table quantization policy. Must satisfy WeightQuantPolicy; defaults to NoWeightQuant.

Constructor & Destructor Documentation

◆ TokenEmbedding()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::TokenEmbedding	(	const std::string &	name,
		const TokenEmbeddingConfig &	config,
		std::optional< DeviceId >	device_id = std::nullopt )

inlineexplicit

Construct a TokenEmbedding component.

Parameters

name	Component name identifier.
config	TokenEmbedding configuration.
device_id	Optional DeviceId for standalone (owned context) mode.

Member Function Documentation

◆ backward()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

TokenIndexType & Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::backward	(	const TokenIndexType &	input,
		const EmbeddingTensorType &	output_grad )

inline

Backward pass – accumulates gradients into wte.

Token indices are discrete and non-differentiable; the returned input_grad tensor exists for interface consistency but carries no meaningful gradient.

wte_grad buffers use atomicAdd accumulation and must be zeroed before each backward call, which zeroGradients() handles.

Parameters

input	Token indices used in forward [B, T].
output_grad	Upstream gradient w.r.t. embeddings [B, T, C].

Returns: Reference to component-owned (unused) input gradient tensor.

Exceptions

std::runtime_error if not built, not in training mode, or buffers are not initialized.

◆ forward()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

EmbeddingTensorType & Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::forward ( const TokenIndexType & input )

inline

Forward pass – returns component-owned embeddings tensor.

output[b, t, :] = wte[ X[b, t], : ]

Accepts any sequence length T <= max built T, including T=1 for single-token autoregressive steps.

Parameters

input Token indices [B, T].

Returns: Reference to component-owned embeddings [B, T, C].

Exceptions

std::runtime_error if the component is not built.

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

DeviceId Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::getDeviceId ( ) const

inlineoverridevirtual

Get the compute device id associated with this component.

Must return the device on which parameters and operations execute.

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

std::vector< ITensor * > Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::getGradients ( ) const

inlineoverridevirtual

Return non-owning pointers to parameter gradient tensors.

Gradient buffers are allocated only when the component is built in training mode, so a component built for inference returns an empty vector. Stateless components return empty in either mode. This is the accessor counterpart to getParameters() and does not throw on mode.

Returns: Vector of gradient pointers; empty when built for inference.

Exceptions

std::runtime_error if called before the component has been built.

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ getMemoryStats()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

MemoryStats Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::getMemoryStats ( ) const

inlineoverridevirtual

Return the current memory allocation breakdown for this component.

Reflects allocations at the moment of the call. The returned stats naturally track the component lifecycle:

After construction – nothing; construction allocates none After build( Inference ) – parameters + T=1 state buffers After build( Training ) – parameters + T=full state buffers After setTrainingMode( Train ) – parameters + state + gradients

For CompositeComponent and Network, the returned stats are the recursive aggregate of all child components.

May be called at any time – no lifecycle preconditions.

Returns: MemoryStats reflecting current allocations.

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ getParameterNames()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

std::vector< std::string > Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::getParameterNames ( ) const

inlineoverridevirtual

List all available parameter names for this component.

Returns an empty vector by default. Leaf components with parameters should override to return their canonical parameter name list in the same stable order used by save_() and loadParameter().

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

std::vector< ITensor * > Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::getParameters ( ) const

inlineoverridevirtual

Return non-owning pointers to parameter tensors.

The returned tensor pointers remain valid for the lifetime of the component. Order should be canonical (weights before biases).

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ getRequiredMemory()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

MemoryStats Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::getRequiredMemory ( const BuildContext & context ) const

inlineoverridevirtual

What onBuilding() would allocate for this context, without allocating.

Mirrors initializeParameters() and onBuilding(). The table dominates every other allocation in this component, which is what makes an unquantized table in a quantized model visible here without running anything. See Specifications/MemoryFootprint.md.

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ getType()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

const ComponentType Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::getType ( ) const

inlineoverridevirtual

Get the component type identifier.

Used for serialization and runtime type identification.

Returns: Component type enum value.

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ loadParameter()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::loadParameter	(	const std::string &	name,
		const ITensorBlob &	blob )

inlineoverridevirtual

Load a parameter from serialized tensor data.

Loads raw tensor bytes directly into an existing parameter tensor, handling precision conversion and device upload as needed.

The component validates that the blob's shape matches the parameter's expected shape, then delegates to the backend to perform:

Precision conversion (blob dtype -> parameter dtype)
Device upload (CPU bytes -> target device)

Takes the parameter name used to locate the target tensor, and a blob holding serialized tensor metadata and raw bytes.

Exceptions

std::runtime_error	if component has no parameters to load.
std::runtime_error	if blob shape doesn't match parameter shape.

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ onBuilding()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::onBuilding ( const BuildContext & build_context )

inlineoverrideprotectedvirtual

Hook invoked by build() to allocate component buffers.

Receives the stored BuildContext. Implementations must use config.allocationSeqLen() when sizing output buffers – this is the single call that makes Inference and Training allocate the correct buffer sizes automatically without per-component logic.

// Example -- Linear component:
shape_t out_shape =
{
    config.batchSize(),
    config.allocationSeqLen(),   // 1 for Inference, T for Training
    config_.getOutputFeatures()
};
output_ = std::make_unique<TensorType>( device, out_shape,
    this->getName() + ".output" );

The default implementation forwards to the legacy onBuilding( const shape_t& ) overload for backwards compatibility. New components should override this overload directly.

Note: Do not call build() or onBuilding() from within this hook.; Implementations should either succeed fully or leave no partial state, as a failed build() may be retried.

Takes the build-time configuration; use its allocationSeqLen() to obtain the correct output buffer sequence dimension.

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::onExecutionContextSet ( )

inlineoverrideprotectedvirtual

Lifecycle hook: Called immediately after ExecutionContext is set.

Override this to perform initialization that requires a valid ExecutionContext. At the time this is called, getExecutionContext() is guaranteed to return a valid context.

Common uses:

Composite components: Create and configure child components.
Device resource allocation: Query device capabilities.

Default implementation does nothing.

Exceptions

Any	exception thrown will cause setExecutionContext() to fail and restore the component to a "context not set" state.

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::onTrainingModeChanging ( TrainingMode training_mode )

inlineoverrideprotectedvirtual

Hook called before TrainingMode transitions.

Called by setTrainingMode() after validation and lock acquisition, before the internal state is updated. Derived classes override to respond to the transition – e.g. zeroing gradient buffers on transition to Eval, or re-enabling dropout on transition to Training.

The default implementation is a no-op.

Takes the incoming TrainingMode.

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

dim_t Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::parameterCount ( ) const

inlineoverridevirtual

Return number of trainable parameters.

For leaf components this is the element count of owned parameter tensors. CompositeComponent and Network implementations should return the recursive aggregate across all children.

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ save_()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::save_	(	ModelArchive &	archive,
		SerializationMode	mode ) const

inlineoverridevirtual

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ saveFlatTensors()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::saveFlatTensors	(	Serialization::SafeTensorsWriter &	writer,
		const std::string &	prefix,
		Serialization::TensorSavePass	pass ) const

inlineoverridevirtual

Emit the table, and its per-row scales when the table is quantized.

The scales ride as a sibling name rather than joining getParameterNames(), for the same reason as Linear: that vector is the archive's save/load join. Underscore, not dot – parseParameterPath() splits on the last dot, so "wte.scales" would resolve to a component named "<prefix>.wte" and the artifact could never be read back.

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::synchronize ( )

inlineoverridevirtual

         @brief Convenience accessor -- true if currently in Eval mode.

         Equivalent to getTrainingMode() == TrainingMode::Eval.
         Valid for both RuntimeMode::Inference and RuntimeMode::Training
         built components.

         @return true if in Eval mode.
        &zwj;/

bool isEvalMode() const noexcept { return getTrainingMode() == TrainingMode::Eval; }

    RuntimeMode getRuntimeMode() const noexcept
    {
        return build_context_.getRuntimeMode();
    }

    bool isInferenceMode() const noexcept
    {
        return build_context_.isInferenceMode();
    }

    bool isTrainingMode() const noexcept
    {
        return build_context_.isTrainingMode();
    }

====================================================================

Synchronization

    /**
       @brief Wait for outstanding device work submitted by this component.

       On CPU this may be a no-op. Use to ensure results are visible to
       the host or to measure synchronous timings.

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ toString()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

std::string Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::toString ( ) const

inlineoverridevirtual

Produce a short, human-readable description of the component.

Implementations should keep output concise and avoid throwing.

Implements Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

◆ zeroGradients()

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

void Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::zeroGradients ( )

inlineoverridevirtual

Clear all model-owned gradients for this component.

Default implementation is a no-op. Composite components should override to recurse to children. Leaf components should override to zero their parameter and activation gradients using device-aware helpers.

Reimplemented from Mila::Dnn::Component< TDeviceType, dtype_t::FP32 >.

Member Data Documentation

◆ kTableDtype

template<DeviceType TDeviceType, TensorDataType TIndex = dtype_t::INT32, TensorDataType TPrecision = dtype_t::FP32, WeightQuantPolicy TTableQuantization = NoWeightQuant>

TensorDataType Mila::Dnn::TokenEmbedding< TDeviceType, TIndex, TPrecision, TTableQuantization >::kTableDtype

staticconstexpr

Initial value:

= kIsQuantized

? TTableQuantization::kStorageDtype : TPrecision

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

Mila/Src/Dnn/Components/Embeddings/TokenEmbedding.ixx

Public Types

Public Member Functions

Static Public Attributes

Protected Member Functions

Additional Inherited Members

Detailed Description

Constructor & Destructor Documentation

◆ TokenEmbedding()

Member Function Documentation

◆ backward()

◆ forward()

◆ getDeviceId()

◆ getGradients()

◆ getMemoryStats()

◆ getParameterNames()

◆ getParameters()

◆ getRequiredMemory()

◆ getType()

◆ loadParameter()

◆ onBuilding()

◆ onExecutionContextSet()

◆ onTrainingModeChanging()

◆ parameterCount()

◆ save_()

◆ saveFlatTensors()

◆ synchronize()

Synchronization

◆ toString()

◆ zeroGradients()

Member Data Documentation

◆ kTableDtype