Stacked experts: gate_up_proj [E, 2I, H] and down_proj [E, H, I], checkpoint-named. More...

Inheritance diagram for Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >:

Public Types
using	ComponentBase = Component<TDeviceType, TPrecision>
using	IndexTensorType = Tensor<TensorDataType::INT32, MR>
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	ScaleTensorType = Tensor<TensorDataType::FP32, MR>
using	TensorType = Tensor<TPrecision, MR>
using	WeightTensorType = Tensor<kWeightDtype, MR>

Public Member Functions
	MixtureOfExperts (const std::string &name, const MixtureOfExpertsConfig &config, std::optional< DeviceId > device_id=std::nullopt)
TensorType &	forward (const TensorType &input, const TensorType &weights, const IndexTensorType &indices)
	Run every token through its selected experts and combine.
DeviceId	getDeviceId () const override
	Get the compute device id associated with this component.
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< ObservableStage >	getObservableStages () const override
	Stages this component publishes, and the passes each is published on.
std::vector< std::string >	getParameterNames () const override
	List all available parameter names for this component.
std::vector< ITensor * >	getParameters () const override
	The stored tensors, scales included, so storage accounting sees every byte.
MemoryStats	getRequiredMemory (const BuildContext &context) const override
	Report what build( context ) would allocate, without allocating it.
const ComponentType	getType () const override
	Get the component type identifier.
void	loadParameter (const std::string &name, const ITensorBlob &blob) override
	Load a parameter from serialized tensor data.
dim_t	parameterCount () const override
	Logical expert weights, whatever their storage.
void	save_ (ModelArchive &archive, SerializationMode) const override
void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const override
	Drive this component's tensors through one pass of a flat safetensors save.
void	synchronize () override
	Wait for outstanding device work submitted by this component.
std::string	toString () const override
	Produce a short, human-readable description of the component.
Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
	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).
IExecutionContext *	getExecutionContext () const
	Get the shared execution context.
const std::string	getName () const
	Get the component's name identifier.
Compute::ComputePassMask	getObservedPasses () const noexcept
virtual std::vector< const ITensor * >	getOutputs () const
	Return non-owning pointers to the tensors this component produces.
RuntimeMode	getRuntimeMode () const noexcept
TrainingMode	getTrainingMode () const noexcept
	The current runtime behavioral mode of this Component.
virtual bool	isBuilt () const final
	Returns true if build() has completed successfully.
bool	isInferenceMode () const noexcept
bool	isTrainingMode () const noexcept
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	setObservedPasses (Compute::ComputePassMask passes) noexcept
	Select which passes this component publishes on.
void	setTrainingMode (TrainingMode mode)
	Set the runtime behavioral mode for this Component.
virtual void	zeroGradients ()
	Clear all model-owned gradients for this component.

Static Public Attributes
static constexpr bool	kIsFp4
static constexpr bool	kIsQuantized = TWeightQuantization::kIsQuantized
static constexpr TensorDataType	kWeightDtype = kIsQuantized ? TWeightQuantization::kStorageDtype : TPrecision

Protected Member Functions
void	onBuilding (const BuildContext &context) override
	Hook invoked by build() to allocate component buffers.
void	onExecutionContextSet () override
	Lifecycle hook: Called immediately after ExecutionContext is set.
void	onTrainingModeChanging (TrainingMode) override
	Hook called before TrainingMode transitions.
Protected Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
bool	hasExecutionContext () const noexcept
	Check if execution context has been set.
template<TensorDataType TParameterPrecision, typename TMemoryResource>
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	publish (Compute::ComputePass pass, std::string_view stage, const ITensor &value) const
	Publish one activation to the observer, if this pass is being observed.
template<TensorDataType TParameterPrecision, typename TMemoryResource>
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>".
template<TensorDataType TParameterPrecision, typename TMemoryResource>
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, TPrecision >
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, TPrecision >
using	HostStagingMemoryResource = typename DeviceTypeTraits<TDeviceType>::host_staging_memory_resource
	Host memory a device-resident parameter stages through.
Protected Attributes inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
BuildContext	build_context_ { shape_t{ 1 }, RuntimeMode::Training }
	The BuildContext stored at build time.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >

Stacked experts: gate_up_proj [E, 2I, H] and down_proj [E, H, I], checkpoint-named.

An expert is a row of these tensors, never an object. Inference-only. Under a per-group FP4 policy each row is stored as a Linear FP4 weight row is – packed nibbles plus per-group FP32 scales, saved as gate_up_proj_scale and down_proj_scale – and a BF16 source is quantized on load.

Template Parameters

TGate	Gate activation: Silu (SwiGLU experts) or Gelu (GeGLU experts, Gemma).
TWeightQuantization	NoWeightQuant or PerGroupFp4<g>; the operation refuses any other.

Member Function Documentation

◆ forward()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

TensorType & Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::forward	(	const TensorType &	input,
		const TensorType &	weights,
		const IndexTensorType &	indices )

inline

Run every token through its selected experts and combine.

Parameters

input	[..., hidden_size], what the experts read.
weights	[..., top_k] combine weights, as Router produces them.
indices	[..., top_k] expert indices, as Router produces them.

Returns: [..., hidden_size], component-owned, valid until the next forward().

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

DeviceId Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

std::vector< ITensor * > Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

◆ getMemoryStats()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

MemoryStats Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

◆ getObservableStages()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

std::vector< ObservableStage > Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::getObservableStages ( ) const

inlineoverridevirtual

Stages this component publishes, and the passes each is published on.

Lets a caller tell an uninstrumented component from a clean one before running anything. The pass qualifier is load-bearing rather than decorative: the stage set genuinely differs by pass, since a prefill can produce intermediates its decode never materializes.

Returns: Vector of stages; empty for components that publish nothing.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ getParameterNames()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

std::vector< std::string > Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

std::vector< ITensor * > Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::getParameters ( ) const

inlineoverridevirtual

The stored tensors, scales included, so storage accounting sees every byte.

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

◆ getRequiredMemory()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

MemoryStats Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::getRequiredMemory ( const BuildContext & context ) const

inlineoverridevirtual

Report what build( context ) would allocate, without allocating it.

The predictive counterpart to getMemoryStats(): what this component would need against what it currently has. Called on a constructed but unbuilt component, which is possible because construction commits no device memory – see Specifications/MemoryFootprint.md.

Implementations mirror their own onBuilding(), which is why the two take the same BuildContext: an implementation cannot answer for a different shape, runtime mode, or prefill size than the one build() would receive.

Composites mirror whatever their getMemoryStats() twin does – children, plus directly-owned buffers, minus any pooling, output-sharing or weight- tying correction. A plain sum over children overcounts wherever a buffer is installed rather than self-allocated.

Parameters

context The context that would be passed to build().

Returns: MemoryStats the component would allocate.

Exceptions

std::logic_error if this component has not implemented the contract.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ getType()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

const ComponentType Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

◆ loadParameter()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

void Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

◆ onBuilding()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

void Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::onBuilding ( const BuildContext & 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, TPrecision >.

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

void Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

void Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::onTrainingModeChanging ( TrainingMode )

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, TPrecision >.

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

dim_t Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::parameterCount ( ) const

inlineoverridevirtual

Logical expert weights, whatever their storage.

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

◆ save_()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

void Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::save_	(	ModelArchive &	archive,
		SerializationMode	) const

inlineoverridevirtual

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

◆ saveFlatTensors()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

void Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::saveFlatTensors	(	Serialization::SafeTensorsWriter &	writer,
		const std::string &	prefix,
		Serialization::TensorSavePass	pass ) const

inlineoverridevirtual

Drive this component's tensors through one pass of a flat safetensors save.

This is the path save_() cannot serve. A quantized weight is packed storage plus a scale companion and ModelArchive has no representation for that pairing, so save_() refuses it; the flat container expresses it as sibling tensors and this writes them.

The vocabulary is the flat format's: a fully qualified dotted component path plus a parameter name, exactly what loadParameters() parses on the way back in. Composites contribute nothing themselves – they recurse and extend the prefix – so the emitted set matches what the reader will route back through loadParameter().

Concrete parameter types, not ITensor, on purpose: a type-erased walk would have to ask at runtime whether a tensor's memory is host-accessible, and MemoryResource::is_host_accessible is a compile-time constant. Keeping concrete types here is what avoids widening an exported core type for a serialization concern.

The default refuses rather than writing nothing. A component that owns parameters and does not implement this would otherwise contribute silently to an artifact that loads, runs, and produces garbage – the same failure mode Phase 0 removed from save_().

Parameters

writer	Writer being driven.
prefix	Fully qualified component path, e.g. "tf_layer_0.qkv_proj".
pass	Declare reserves byte ranges; Write streams bytes in declaration order.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

void Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::synchronize ( )

inlineoverridevirtual

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, TPrecision >.

◆ toString()

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

std::string Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::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, TPrecision >.

Member Data Documentation

◆ kIsFp4

template<DeviceType TDeviceType, TensorDataType TPrecision, ActivationType TGate = ActivationType::Silu, WeightQuantPolicy TWeightQuantization = NoWeightQuant>

bool Mila::Dnn::MixtureOfExperts< TDeviceType, TPrecision, TGate, TWeightQuantization >::kIsFp4

staticconstexpr

Initial value:

                             = []() constexpr
{
    if constexpr ( requires { TWeightQuantization::kIsFp4E2M1; } )
    {
        return TWeightQuantization::kIsQuantized && TWeightQuantization::kIsFp4E2M1;
    }
    else
    {
        return false;
    }
}()

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

Mila/Src/Dnn/Components/MixtureOfExperts/MixtureOfExperts.ixx

Public Types

Public Member Functions

Static Public Attributes

Protected Member Functions

Additional Inherited Members

Detailed Description

Member Function Documentation

◆ forward()

◆ getDeviceId()

◆ getGradients()

◆ getMemoryStats()

◆ getObservableStages()

◆ getParameterNames()

◆ getParameters()

◆ getRequiredMemory()

◆ getType()

◆ loadParameter()

◆ onBuilding()

◆ onExecutionContextSet()

◆ onTrainingModeChanging()

◆ parameterCount()

◆ save_()

◆ saveFlatTensors()

◆ synchronize()

◆ toString()

Member Data Documentation

◆ kIsFp4