|
Mila
Deep Neural Network Library
|
A component that contains and manages child components. More...
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 ¶m_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 ¶meter_name, const Tensor< TParameterPrecision, TMemoryResource > ¶meter) 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 > ¶meter, 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. | |
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:
NOTE:
|
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).
| name | Component name identifier (mandatory) |
| std::invalid_argument | if name is not a valid identifier |
|
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:
| component | Shared pointer to the constructed component |
| std::runtime_error | if called after build() |
| std::invalid_argument | if component is null |
| std::invalid_argument | if component name already exists |
| std::invalid_argument | if component already has its own ExecutionContext |
|
inlinenoexcept |
Get the number of direct children.
|
inline |
Clear all child components.
| std::runtime_error | if called after build() |
|
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.
|
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.
|
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.
| name | Name of the direct child component |
| std::out_of_range | if the direct child is not found |
|
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.
| TComponent | Expected component type |
| name | Name of the direct child component |
| std::out_of_range | if component name not found |
| std::runtime_error | if dynamic cast fails (type mismatch) |
|
inline |
Get all child components in insertion order.
|
inlineoverridevirtual |
Get the compute device for this composite.
Returns the device from the shared execution context.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
Reimplemented in Mila::Dnn::Network< TDeviceType, TPrecision >.
|
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().
| std::runtime_error | if called before build() |
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlineoverridevirtual |
Get all parameters from all children.
| std::runtime_error | if called before build() |
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
|
inlinenoexcept |
Check if this composite has any children.
|
inline |
Check if a named child component exists.
| name | Name to query |
|
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.
| archive | Archive to read from, scoped to this composite. |
| mode | Serialization mode (passed to children). |
| std::runtime_error | if 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 >.
|
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.
|
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 >.
|
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().
| training_mode | New 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 >.
|
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.
|
inlineoverridevirtual |
Count parameters across all children.
| std::runtime_error | if 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 >.
|
inline |
Remove a child component by name.
| name | Name of the component to remove |
| std::runtime_error | if called after build() |
|
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 >.
|
inlineoverrideprotectedvirtual |
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
Reimplemented in Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GptTransformer< TDeviceType, TPrecision >, Mila::Dnn::LlamaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::Network< TDeviceType, TPrecision >, and Mila::Dnn::QwenTransformer< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >.
|
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 >.
|
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.
|
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 >.
|
inlineoverridevirtual |
Generate a human-readable description.
Implements Mila::Dnn::Component< TDeviceType, TPrecision >.
Reimplemented in Mila::Dnn::GatedMLP< TDeviceType, TPrecision, TGate, TWeightQuantization >, Mila::Dnn::GatedMLP< TDeviceType, TPrecision, ActivationType::Gelu, TWeightQuant >, Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GptBlock< TDeviceType, TPrecision >, Mila::Dnn::GptTransformer< TDeviceType, TPrecision >, Mila::Dnn::LlamaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::MLP< TDeviceType, TPrecision >, Mila::Dnn::Network< 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 >.
|
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");
| path | Dot-separated path (e.g. "layer_0.mlp.fc_1") |