Class FSDPModel<T, TInput, TOutput>
- Namespace
- AiDotNet.DistributedTraining
- Assembly
- AiDotNet.dll
Implements FSDP (Fully Sharded Data Parallel) model wrapper that shards parameters across multiple processes.
public class FSDPModel<T, TInput, TOutput> : ShardedModelBase<T, TInput, TOutput>, IShardedModel<T, TInput, TOutput>, IFullModel<T, TInput, TOutput>, IModel<TInput, TOutput, ModelMetadata<T>>, IModelSerializer, ICheckpointableModel, IParameterizable<T, TInput, TOutput>, IFeatureAware, IFeatureImportance<T>, ICloneable<IFullModel<T, TInput, TOutput>>, IGradientComputable<T, TInput, TOutput>, IJitCompilable<T>
Type Parameters
TThe numeric type
TInputThe input type for the model
TOutputThe output type for the model
- Inheritance
-
ShardedModelBase<T, TInput, TOutput>FSDPModel<T, TInput, TOutput>
- Implements
-
IShardedModel<T, TInput, TOutput>IFullModel<T, TInput, TOutput>IModel<TInput, TOutput, ModelMetadata<T>>IParameterizable<T, TInput, TOutput>ICloneable<IFullModel<T, TInput, TOutput>>IGradientComputable<T, TInput, TOutput>
- Derived
- Inherited Members
- Extension Methods
Remarks
Strategy Overview: FSDP (Fully Sharded Data Parallel) is PyTorch's implementation of the ZeRO-3 optimization strategy. It shards model parameters, gradients, and optimizer states across all processes, achieving maximum memory efficiency. Parameters are gathered just-in-time for forward/backward passes and then released.
For Beginners: This class implements FSDP (Fully Sharded Data Parallel), which makes any model work across multiple GPUs or machines with maximum memory efficiency. It automatically handles: - Splitting ALL model components (parameters, gradients, optimizer states) across processes - Gathering parameters only when needed for forward/backward pass - Releasing parameters immediately after use to save memory - Averaging gradients across all processes during training
Think of it like a team project where each person holds part of the solution, but unlike DDP, FSDP only shares the full model temporarily when absolutely needed, then immediately goes back to holding just their piece. This saves a lot of memory!
Use Cases: - Training very large models that don't fit in a single GPU's memory - Maximizing memory efficiency for multi-GPU training - Scaling to hundreds or thousands of GPUs
Trade-offs: - Memory: Excellent - shards everything (parameters + gradients + optimizer states) - Communication: Higher - requires AllGather for each forward/backward pass - Complexity: Moderate - automatic just-in-time parameter gathering - Best for: Very large models, memory-constrained scenarios
Example:
// Original model
var model = new NeuralNetworkModel<double>(...);
// Wrap it for FSDP distributed training
var backend = new InMemoryCommunicationBackend<double>(rank: 0, worldSize: 4);
var config = new ShardingConfiguration<double>(backend);
var fsdpModel = new FSDPModel<double, Tensor<double>, Tensor<double>>(model, config);
// Now train as usual - FSDP magic happens automatically!
fsdpModel.Train(inputs, outputs);
Constructors
FSDPModel(IFullModel<T, TInput, TOutput>, IShardingConfiguration<T>)
Creates a new FSDP model wrapping an existing model.
public FSDPModel(IFullModel<T, TInput, TOutput> wrappedModel, IShardingConfiguration<T> config)
Parameters
wrappedModelIFullModel<T, TInput, TOutput>The model to wrap with FSDP capabilities
configIShardingConfiguration<T>Configuration for sharding and communication
Remarks
For Beginners: This constructor takes your existing model and makes it distributed using FSDP strategy. You provide: 1. The model you want to make distributed 2. A configuration that tells us how to do the distribution
The constructor automatically: - Splits the model's parameters across all processes (sharding) - Sets up communication channels - Prepares everything for FSDP distributed training
Exceptions
- ArgumentNullException
Thrown if model or config is null
Methods
Clone()
Creates a shallow copy of this object.
public override IFullModel<T, TInput, TOutput> Clone()
Returns
- IFullModel<T, TInput, TOutput>
DeepCopy()
Creates a deep copy of this object.
public override IFullModel<T, TInput, TOutput> DeepCopy()
Returns
- IFullModel<T, TInput, TOutput>
Deserialize(byte[])
Loads a previously serialized model from binary data.
public override void Deserialize(byte[] data)
Parameters
databyte[]The byte array containing the serialized model data.
Remarks
This method takes binary data created by the Serialize method and uses it to restore a model to its previous state.
For Beginners: This is like opening a saved file to continue your work.
When you call this method:
- You provide the binary data (bytes) that was previously created by Serialize
- The model rebuilds itself using this data
- After deserializing, the model is exactly as it was when serialized
- It's ready to make predictions without needing to be trained again
For example:
- You download a pre-trained model file for detecting spam emails
- You deserialize this file into your application
- Immediately, your application can detect spam without any training
- The model has all the knowledge that was built into it by its original creator
This is particularly useful when:
- You want to use a model that took days to train
- You need to deploy the same model across multiple devices
- You're creating an application that non-technical users will use
Think of it like installing the brain of a trained expert directly into your application.
GetActiveFeatureIndices()
Gets the indices of features that are actively used by this model.
public override IEnumerable<int> GetActiveFeatureIndices()
Returns
GetFeatureImportance()
Gets the feature importance scores.
public override Dictionary<string, T> GetFeatureImportance()
Returns
- Dictionary<string, T>
GetModelMetadata()
Retrieves metadata and performance metrics about the trained model.
public override ModelMetadata<T> GetModelMetadata()
Returns
- ModelMetadata<T>
An object containing metadata and performance metrics about the trained model.
Remarks
This method provides information about the model's structure, parameters, and performance metrics.
For Beginners: Model metadata is like a report card for your machine learning model.
Just as a report card shows how well a student is performing in different subjects, model metadata shows how well your model is performing and provides details about its structure.
This information typically includes:
- Accuracy measures: How well does the model's predictions match actual values?
- Error metrics: How far off are the model's predictions on average?
- Model parameters: What patterns did the model learn from the data?
- Training information: How long did training take? How many iterations were needed?
For example, in a house price prediction model, metadata might include:
- Average prediction error (e.g., off by $15,000 on average)
- How strongly each feature (bedrooms, location) influences the prediction
- How well the model fits the training data
This information helps you understand your model's strengths and weaknesses, and decide if it's ready to use or needs more training.
IsFeatureUsed(int)
Checks if a specific feature is used by this model.
public override bool IsFeatureUsed(int featureIndex)
Parameters
featureIndexint
Returns
LoadModel(string)
Loads the model from a file.
public override void LoadModel(string filePath)
Parameters
filePathstringThe path to the file containing the saved model.
Remarks
This method provides a convenient way to load a model directly from disk. It combines file I/O operations with deserialization.
For Beginners: This is like clicking "Open" in a document editor. Instead of manually reading from a file and then calling Deserialize(), this method does both steps for you.
Exceptions
- FileNotFoundException
Thrown when the specified file does not exist.
- IOException
Thrown when an I/O error occurs while reading from the file or when the file contains corrupted or invalid model data.
Predict(TInput)
Uses the trained model to make predictions for new input data.
public override TOutput Predict(TInput input)
Parameters
inputTInputA matrix where each row represents a new example to predict and each column represents a feature.
Returns
- TOutput
A vector containing the predicted values for each input example.
Remarks
After training, this method applies the learned patterns to new data to predict outcomes.
For Beginners: Prediction is when the model uses what it learned to make educated guesses about new information.
Continuing the fruit identification example:
- After learning from many examples, the child (model) can now identify new fruits they haven't seen before
- They look at the color, shape, and size to make their best guess
In machine learning:
- You give the model new data it hasn't seen during training
- The model applies the patterns it learned to make predictions
- The output is the model's best estimate based on its training
For example, in a house price prediction model:
- You provide features of a new house (square footage, bedrooms, location)
- The model predicts what price that house might sell for
This method is used after training is complete, when you want to apply your model to real-world data.
SaveModel(string)
Saves the model to a file.
public override void SaveModel(string filePath)
Parameters
filePathstringThe path where the model should be saved.
Remarks
This method provides a convenient way to save the model directly to disk. It combines serialization with file I/O operations.
For Beginners: This is like clicking "Save As" in a document editor. Instead of manually calling Serialize() and then writing to a file, this method does both steps for you.
Exceptions
- IOException
Thrown when an I/O error occurs while writing to the file.
- UnauthorizedAccessException
Thrown when the caller does not have the required permission to write to the specified file path.
Serialize()
Converts the current state of a machine learning model into a binary format.
public override byte[] Serialize()
Returns
- byte[]
A byte array containing the serialized model data.
Remarks
This method captures all the essential information about a trained model and converts it into a sequence of bytes that can be stored or transmitted.
For Beginners: This is like exporting your work to a file.
When you call this method:
- The model's current state (all its learned patterns and parameters) is captured
- This information is converted into a compact binary format (bytes)
- You can then save these bytes to a file, database, or send them over a network
For example:
- After training a model to recognize cats vs. dogs in images
- You can serialize the model to save all its learned knowledge
- Later, you can use this saved data to recreate the model exactly as it was
- The recreated model will make the same predictions as the original
Think of it like taking a snapshot of your model's brain at a specific moment in time.
SetActiveFeatureIndices(IEnumerable<int>)
Sets the active feature indices for this model.
public override void SetActiveFeatureIndices(IEnumerable<int> featureIndices)
Parameters
featureIndicesIEnumerable<int>
SynchronizeGradients()
Synchronizes gradients across all processes using AllReduce.
public override void SynchronizeGradients()
Remarks
After this operation, all processes have the same averaged gradients.
For Beginners: During training, each process calculates gradients based on its portion of the data. This method combines (averages) those gradients so that everyone is learning from everyone else's experiences. It's like a team meeting where everyone shares what they learned.
Train(TInput, TOutput)
Trains the model using input features and their corresponding target values.
public override void Train(TInput input, TOutput expectedOutput)
Parameters
inputTInputexpectedOutputTOutput
Remarks
This method takes training data and adjusts the model's internal parameters to learn patterns in the data.
For Beginners: Training is like teaching the model by showing it examples.
Imagine teaching a child to identify fruits:
- You show them many examples of apples, oranges, and bananas (input features x)
- You tell them the correct name for each fruit (target values y)
- Over time, they learn to recognize the patterns that distinguish each fruit
In machine learning:
- The x parameter contains features (characteristics) of your data
- The y parameter contains the correct answers you want the model to learn
- During training, the model adjusts its internal calculations to get better at predicting y from x
For example, in a house price prediction model:
- x would contain features like square footage, number of bedrooms, location
- y would contain the actual sale prices of those houses
WithParameters(Vector<T>)
Creates a new instance with the specified parameters.
public override IFullModel<T, TInput, TOutput> WithParameters(Vector<T> parameters)
Parameters
parametersVector<T>
Returns
- IFullModel<T, TInput, TOutput>