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

A component that contains and manages child components. More...

Inheritance diagram for Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >:
Mila::Dnn::Component< TDeviceType, TPrecision > Mila::Dnn::GatedMLP< TDeviceType, TPrecision, ActivationType::Gelu, TWeightQuant > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts > Mila::Dnn::LlamaBlock< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy > Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy > Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan > Mila::Dnn::GatedMLP< TDeviceType, TPrecision, TGate, TWeightQuantization > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts > Mila::Dnn::GptBlock< TDeviceType, TPrecision > Mila::Dnn::LlamaBlock< TDeviceType, TPrecision, TWeightQuant, TKvPolicy > Mila::Dnn::MLP< TDeviceType, TPrecision > Mila::Dnn::Network< TDeviceType, TPrecision > Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy > Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan > Mila::Dnn::Router< TDeviceType, TPrecision >

Public Types

using ComponentBase = Component<TDeviceType, TPrecision>
using ComponentPtr = std::shared_ptr<Component<TDeviceType, TPrecision>>

Public Member Functions

 CompositeComponent (CompositeComponent &&) noexcept=default
 CompositeComponent (const CompositeComponent &)=delete
 CompositeComponent (const std::string &name)
 Construct composite component with name.
CompositeComponent & addComponent (ComponentPtr component)
 Add a pre-constructed child component (chainable).
size_t childCount () const noexcept
 Get the number of direct children.
void clearComponents ()
 Clear all child components.
std::vector< std::string > componentPaths () const
 Every component path in this subtree, in traversal order.
ComponentPtr findComponent (const std::string &path) const
 Resolve a dot-separated component path within this composite.
ComponentPtr getComponent (const std::string &name) const
 Retrieve a direct child component by name.
const std::vector< ComponentPtr > & getComponents () const
 Get all child components in insertion order.
DeviceId getDeviceId () const override
 Get the compute device for this composite.
std::vector< ITensor * > getGradients () const override
 Get all parameter gradients from all children.
std::vector< ITensor * > getParameters () const override
 Get all parameters from all children.
bool hasChildren () const noexcept
 Check if this composite has any children.
bool hasComponent (const std::string &name) const
 Check if a named child component exists.
size_t observe (std::string_view pattern, Compute::ComputePassMask passes, Compute::ActivationObserver sink)
 Attach sink to every component in this subtree whose path matches.
CompositeComponent & operator= (CompositeComponent &&) noexcept=default
CompositeComponent & operator= (const CompositeComponent &)=delete
dim_t parameterCount () const override
 Count parameters across all children.
bool removeComponent (const std::string &name)
 Remove a child component by name.
void saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const override
 Recurse into children, extending the flat dotted prefix.
void stopObserving ()
 Detach every observer in this subtree and clear the shared sink.
void synchronize () override
 Synchronize all child components.
std::string toString () const override
 Generate a human-readable description.
ComponentPtr tryFindComponent (const std::string &path) const
 Try to resolve a dot-separated component path within this composite.
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.
virtual MemoryStats getMemoryStats () const =0
 Return the current memory allocation breakdown for this component.
const std::string getName () const
 Get the component's name identifier.
virtual std::vector< ObservableStage > getObservableStages () const
 Stages this component publishes, and the passes each is published on.
Compute::ComputePassMask getObservedPasses () const noexcept
virtual std::vector< const ITensor * > getOutputs () const
 Return non-owning pointers to the tensors this component produces.
virtual std::vector< std::string > getParameterNames () const
 List all available parameter names for this component.
virtual MemoryStats getRequiredMemory (const BuildContext &context) const
 Report what build( context ) would allocate, without allocating it.
RuntimeMode getRuntimeMode () const noexcept
TrainingMode getTrainingMode () const noexcept
 The current runtime behavioral mode of this Component.
virtual const ComponentType getType () const =0
 Get the component type identifier.
virtual bool isBuilt () const final
 Returns true if build() has completed successfully.
bool isInferenceMode () const noexcept
bool isTrainingMode () const noexcept
virtual void loadParameter (const std::string &, const Serialization::ITensorBlob &)
 Load a parameter from serialized tensor data.
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.

Protected Member Functions

template<typename TComponent>
std::shared_ptr< TComponent > getComponentAs (const std::string &name) const
 Retrieve a typed child component by name.
void load_ (ModelArchive &archive, SerializationMode mode) override
 Restore children from their nested scopes, mirroring save_().
void onExecutionContextSet () override
 Hook invoked after ExecutionContext is set.
void onTrainingModeChanging (TrainingMode training_mode) override
 Hook invoked when training mode is about to change.
virtual void optimize ()
 Virtual hook for graph optimization after construction.
void requireSerializableParameters () const override
 No-op override: a composite names no parameters of its own.
void save_ (ModelArchive &archive, SerializationMode mode) const override
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.
virtual void onBuilding (const BuildContext &)
 Hook invoked by build() to allocate component buffers.
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.

Friends

class ComponentFactory

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>
class Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >

A component that contains and manages child components.

CompositeComponent is a device-parameterized abstract container that manages child component lifecycle, aggregates operations (parameters, gradients, training mode), and provides context propagation. Derived types implement execution semantics (forward/backward) and architecture definition (createGraph()).

Architecture Philosophy:

  • Context-independent graph creation: Architecture defined without device knowledge
  • Three-phase lifecycle: Graph creation -> Context binding -> Shape binding
  • Automatic context propagation: Base class propagates context to all children
  • Component-owns-name: Children manage their own identity via getName()

NOTE:

Constructor & Destructor Documentation

◆ CompositeComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>
Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::CompositeComponent ( const std::string & name)
inlineexplicit

Construct composite component with name.

The composite component is named for identification in hierarchical structures. Derived classes should call createGraph() from their constructor to define the architecture graph (context-independent).

All child components added via addComponent() will receive ExecutionContext automatically when the composite receives its context (via onExecutionContextSet).

Parameters
nameComponent name identifier (mandatory)
Exceptions
std::invalid_argumentif name is not a valid identifier

Member Function Documentation

◆ addComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>
CompositeComponent & Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::addComponent ( ComponentPtr component)
inline

Add a pre-constructed child component (chainable).

Registers a component that was constructed externally (typically by the derived class in its createGraph() method). The component's getName() is used as the lookup key.

Components are expected to be created in shared mode (no ExecutionContext). Context will be automatically propagated to all children when this composite receives its context via onExecutionContextSet().

Usage pattern in derived class:

void MLP::createGraph()
{
auto fc1 = std::make_shared<LinearType>(config, std::nullopt);
fc1->setName(this->getName() + ".fc1");
this->addComponent(fc1);
// ... more components
}
const std::string getName() const
Get the component's name identifier.
Definition Component.ixx:522
CompositeComponent & addComponent(ComponentPtr component)
Add a pre-constructed child component (chainable).
Definition CompositeComponent.ixx:132
Parameters
componentShared pointer to the constructed component
Returns
Reference to *this for method chaining
Exceptions
std::runtime_errorif called after build()
std::invalid_argumentif component is null
std::invalid_argumentif component name already exists
std::invalid_argumentif component already has its own ExecutionContext

◆ childCount()

template<DeviceType TDeviceType, TensorDataType TPrecision>
size_t Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::childCount ( ) const
inlinenoexcept

Get the number of direct children.

Returns
Number of child components

◆ clearComponents()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::clearComponents ( )
inline

Clear all child components.

Exceptions
std::runtime_errorif called after build()

◆ componentPaths()

template<DeviceType TDeviceType, TensorDataType TPrecision>
std::vector< std::string > Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::componentPaths ( ) const
inline

Every component path in this subtree, in traversal order.

The structural view: what is there to observe, before deciding what to observe. Also what makes a zero from observe() diagnosable rather than merely disappointing.

◆ findComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>
ComponentPtr Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::findComponent ( const std::string & path) const
inline

Resolve a dot-separated component path within this composite.

Supports both relative paths ("lenc.wte") and absolute paths ("gpt2.lenc.wte"). If path starts with this component's name, strips it before searching.

◆ getComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>
ComponentPtr Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getComponent ( const std::string & name) const
inline

Retrieve a direct child component by name.

This performs direct (non-recursive) lookup of immediate children only. Use findComponent() to resolve dot-separated paths across the subgraph.

Parameters
nameName of the direct child component
Returns
Shared pointer to the component
Exceptions
std::out_of_rangeif the direct child is not found

◆ getComponentAs()

template<DeviceType TDeviceType, TensorDataType TPrecision>
template<typename TComponent>
std::shared_ptr< TComponent > Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getComponentAs ( const std::string & name) const
inlineprotected

Retrieve a typed child component by name.

Helper method for derived composites (like MLP) that need to cache typed pointers to children in their onBuilding() hook. Performs dynamic_pointer_cast and validates the cast succeeded.

Note: This resolves direct children only. For full-path resolution use findComponent() on the appropriate root composite or Network.

Template Parameters
TComponentExpected component type
Parameters
nameName of the direct child component
Returns
Shared pointer to component with correct type
Exceptions
std::out_of_rangeif component name not found
std::runtime_errorif dynamic cast fails (type mismatch)

◆ getComponents()

template<DeviceType TDeviceType, TensorDataType TPrecision>
const std::vector< ComponentPtr > & Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getComponents ( ) const
inline

Get all child components in insertion order.

Returns
Vector of child component pointers

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TPrecision>
DeviceId Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getDeviceId ( ) const
inlineoverridevirtual

Get the compute device for this composite.

Returns the device from the shared execution context.

Returns
DeviceId for this composite and its children

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

Reimplemented in Mila::Dnn::Network< TDeviceType, TPrecision >.

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TPrecision>
std::vector< ITensor * > Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getGradients ( ) const
inlineoverridevirtual

Get all parameter gradients from all children.

Gradient buffers are allocated only when built in training mode, so a component built for inference yields an empty result. This mirrors the leaf-component contract (e.g. Linear) and the symmetry with getParameters().

Returns
Vector of non-owning pointers to gradient tensors (empty in inference mode)
Exceptions
std::runtime_errorif called before build()

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

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision>
std::vector< ITensor * > Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getParameters ( ) const
inlineoverridevirtual

Get all parameters from all children.

Returns
Vector of non-owning pointers to parameter tensors
Exceptions
std::runtime_errorif called before build()

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

◆ hasChildren()

template<DeviceType TDeviceType, TensorDataType TPrecision>
bool Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::hasChildren ( ) const
inlinenoexcept

Check if this composite has any children.

Returns
true if at least one child component exists

◆ hasComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>
bool Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::hasComponent ( const std::string & name) const
inline

Check if a named child component exists.

Parameters
nameName to query
Returns
true if a child with this name exists

◆ load_()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::load_ ( ModelArchive & archive,
SerializationMode mode )
inlineoverrideprotectedvirtual

Restore children from their nested scopes, mirroring save_().

Walks the same children in the same order under the same scopes, so a child's blobs are read from exactly the paths they were written to. Restores into the live graph – the composite and its children must already be constructed and built.

Parameters
archiveArchive to read from, scoped to this composite.
modeSerialization mode (passed to children).
Exceptions
std::runtime_errorif the composite is not built.

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

Reimplemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GptTransformer< TDeviceType, TPrecision >, Mila::Dnn::Network< TDeviceType, TPrecision >, and Mila::Dnn::Router< TDeviceType, TPrecision >.

◆ observe()

template<DeviceType TDeviceType, TensorDataType TPrecision>
size_t Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::observe ( std::string_view pattern,
Compute::ComputePassMask passes,
Compute::ActivationObserver sink )
inline

Attach sink to every component in this subtree whose path matches.

The door onto the publication machinery. Every component publishes its activations on every inference pass, but publication is gated per component by setObservedPasses and routed through one observer on the execution context – so without this walk a caller has to reach each component by hand, which is what pushed the first consumer into bolting a narrow accessor onto the model instead.

pattern is a component path with * matching any run of characters, dots included – so a * reaches every depth below where it stands:

"qwen.lm_head"   one component
"qwen.blk_3.*"   everything inside block 3, but not the block itself
"qwen.blk_*"     every block and everything inside them
"*"              the whole subtree

No pattern selects the blocks alone. Attach to the wider set and route inside the sink on its path argument.

Resolved ONCE, here. Matching cost lands at attachment and never on a publication, which is what makes a whole-tree pattern affordable during a real run.

The sink is per execution context, so the most recent one wins – that is the context's design, not a limitation of this walk. Attach once with one sink and route inside it on the path argument; a second call with a different sink silently redirects the first pattern's components too. Passes, by contrast, accumulate: two calls with different patterns leave both sets of components observing.

Returns
How many components matched. Zero means nothing will ever publish – check it. A pattern that matches nothing is indistinguishable, downstream, from a model that produced nothing worth reporting, and that false negative is exactly what a NaN hunt must not have.

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::onExecutionContextSet ( )
inlineoverrideprotectedvirtual

Hook invoked after ExecutionContext is set.

Propagates the execution context to all child components that don't already have one. This enables the pattern where composites define their architecture graph in the constructor (context-independent) and context is bound later when available.

Called by Component::setExecutionContext() after the context is registered. Automatically invoked for both standalone mode (component creates own context) and shared mode (parent provides context).

Override this in derived classes if additional context-dependent initialization is required beyond context propagation to children.

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

Reimplemented in Mila::Dnn::Router< TDeviceType, TPrecision >.

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::onTrainingModeChanging ( TrainingMode training_mode)
inlineoverrideprotectedvirtual

Hook invoked when training mode is about to change.

Propagates the new mode to all child components. The hook runs with the Component's training mutex held; it MUST NOT call setTrainingMode().

Parameters
training_modeNew training mode (Normal or Eval)

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

Reimplemented in Mila::Dnn::GatedMLP< TDeviceType, TPrecision, TGate, TWeightQuantization >, Mila::Dnn::GatedMLP< TDeviceType, TPrecision, ActivationType::Gelu, TWeightQuant >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GptBlock< TDeviceType, TPrecision >, Mila::Dnn::GptTransformer< TDeviceType, TPrecision >, Mila::Dnn::LlamaBlock< TDeviceType, TPrecision, TWeightQuant, TKvPolicy >, Mila::Dnn::LlamaBlock< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::LlamaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::MLP< TDeviceType, TPrecision >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, Mila::Dnn::QwenTransformer< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, and Mila::Dnn::Router< TDeviceType, TPrecision >.

◆ optimize()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::optimize ( )
inlineprotectedvirtual

Virtual hook for graph optimization after construction.

Called automatically after createGraph() completes. Derived classes can override to perform fusion, pruning, or other optimizations.

Default implementation does nothing. Override to perform architecture-specific graph optimizations.

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TPrecision>
dim_t Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::parameterCount ( ) const
inlineoverridevirtual

Count parameters across all children.

Returns
Total number of trainable parameters
Exceptions
std::runtime_errorif called before build()

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

Reimplemented in Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, and Mila::Dnn::Router< TDeviceType, TPrecision >.

◆ removeComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>
bool Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::removeComponent ( const std::string & name)
inline

Remove a child component by name.

Parameters
nameName of the component to remove
Returns
true if removed, false if not found
Exceptions
std::runtime_errorif called after build()

◆ requireSerializableParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::requireSerializableParameters ( ) const
inlineoverrideprotectedvirtual

No-op override: a composite names no parameters of its own.

parameterCount() on a composite sums its children, so the base implementation would demand names this component never owns. Each child is checked individually by the recursion in save_().

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

◆ save_()

◆ saveFlatTensors()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::saveFlatTensors ( Serialization::SafeTensorsWriter & writer,
const std::string & prefix,
Serialization::TensorSavePass pass ) const
inlineoverridevirtual

Recurse into children, extending the flat dotted prefix.

A composite contributes no tensors of its own. It exists here to turn the component tree into the flat vocabulary the weights format uses – the same dotted paths loadParameters() splits with parseParameterPath() and resolves with findComponent(), so what this writes is exactly what that reads.

Children carry fully qualified names ("gemma.tf_layer_0.qkv_proj"), so the relative segment is recovered by stripping this composite's own name. Calling the root with an empty prefix therefore drops the model name and yields "tf_layer_0.qkv_proj", matching the converter's naming rather than a Mila-internal path.

child_components_ rather than child_component_map_: the map is unordered, and the writer requires bodies in declaration order, so an unordered walk would produce a file whose data region disagrees with its own index.

Public, unlike save_() which is public on Component but protected here – an asymmetry already filed as a defect. A flat save is driven from outside the tree, so repeating it would force every caller through a forwarder.

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

Reimplemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, and Mila::Dnn::Router< TDeviceType, TPrecision >.

◆ stopObserving()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::stopObserving ( )
inline

Detach every observer in this subtree and clear the shared sink.

Necessary rather than tidy: an observer left attached keeps firing on every pass of every subsequent run, and a probe that outlives its question is both a cost and a source of confusing output.

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TPrecision>
void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::synchronize ( )
inlineoverridevirtual

Synchronize all child components.

Waits for outstanding device operations on all children.

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

Reimplemented in Mila::Dnn::Network< TDeviceType, TPrecision >.

◆ toString()

◆ tryFindComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>
ComponentPtr Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::tryFindComponent ( const std::string & path) const
inline

Try to resolve a dot-separated component path within this composite.

Non-throwing version: returns nullptr if any segment is not found or if a path segment attempts to traverse into a non-composite leaf.

Example: auto ptr = composite->tryFindComponent("encoder.mlp.fc1");

Parameters
pathDot-separated path (e.g. "layer_0.mlp.fc_1")
Returns
ComponentPtr or nullptr when not found / invalid traversal

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