CUDA implementation of the fused softmax and cross entropy operation for neural networks. More...

Inheritance diagram for Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >:

[legend]

Collaboration diagram for Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >:

[legend]

Public Types
using	MR = typename CudaDevice::MR

Public Types inherited from Mila::Dnn::Compute::BinaryOperation< TPrecision, int, TPrecision, DeviceType::Cuda >
using	MR = std::conditional_t< TDeviceType==DeviceType::Cuda, CudaMemoryResource, HostMemoryResource >
	Memory resource type based on device type.

Public Member Functions
	FusedSoftmaxCrossEntropyOp ()
	Constructs a new CUDA Fused Softmax Cross Entropy operation with the default device context.

	FusedSoftmaxCrossEntropyOp (std::shared_ptr< DeviceContext > context)
	Constructs a new CUDA Fused Softmax Cross Entropy operation with a specific device context.

void	backward (const Tensor< TPrecision, MR > &input1, const Tensor< int, MR > &input2, const Tensor< TPrecision, MR > &output, const Tensor< TPrecision, MR > &output_gradient, const std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &parameters, std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &parameter_gradients, Tensor< TPrecision, MR > &input1_gradient, Tensor< int, MR > &input2_gradient, const OperationAttributes &properties, const std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &output_state) const override
	Performs the backward pass of the fused softmax cross entropy operation.

void	forward (const Tensor< TPrecision, MR > &logits, const Tensor< int, MR > &targets, const std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &parameters, const OperationAttributes &properties, Tensor< TPrecision, MR > &losses, std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &output_state) const override
	Performs the forward pass of the fused softmax cross entropy operation on CUDA.

std::string	getName () const override
	Gets the name of this operation.

Public Member Functions inherited from Mila::Dnn::Compute::BinaryOperation< TPrecision, int, TPrecision, DeviceType::Cuda >
	BinaryOperation (OperationType operation_type)
	Constructs a BinaryOperation with the specified operation type and precision policy.

	BinaryOperation (OperationType operation_type, std::shared_ptr< DeviceContext > context)
	Constructs a BinaryOperation with the specified operation type, device context, and precision policy.

virtual	~BinaryOperation ()=default
	Virtual destructor for proper cleanup of derived classes.

virtual void	backward (const Tensor< int, MR > &input1, const Tensor< TPrecision, MR > &input2, const Tensor< DeviceType::Cuda, MR > &output, const Tensor< DeviceType::Cuda, MR > &output_gradient, const std::vector< std::shared_ptr< Tensor< int, MR > > > &parameters, std::vector< std::shared_ptr< Tensor< DeviceType::Cuda, MR > > > &parameter_gradients, Tensor< int, MR > &input1_gradient, Tensor< TPrecision, MR > &input2_gradient, const OperationAttributes &attributes, const std::vector< std::shared_ptr< Tensor< DeviceType::Cuda, MR > > > &output_state) const
	Executes the backward pass of a binary operation.

virtual void	forward (const Tensor< int, MR > &input1, const Tensor< TPrecision, MR > &input2, const std::vector< std::shared_ptr< Tensor< int, MR > > > &parameters, const OperationAttributes &attributes, Tensor< DeviceType::Cuda, MR > &output, std::vector< std::shared_ptr< Tensor< DeviceType::Cuda, MR > > > &output_state) const=0
	Executes the forward pass of a binary operation.

Public Member Functions inherited from Mila::Dnn::Compute::OperationBase< TDeviceType, TInput1, TInput2, TOutput >
	OperationBase (OperationType operation_type, std::shared_ptr< DeviceContext > context)
	Constructs an OperationBase object with a specific device context and compute precision.

virtual	~OperationBase ()=default
	Virtual destructor for the OperationBase class.

std::shared_ptr< DeviceContext >	getDeviceContext () const
	Gets the device context associated with this operation.

DeviceType	getDeviceType () const
	Gets the device type for this operation.

OperationType	getOperationType () const
	Gets the operation type enumeration value.

Detailed Description

template<typename TPrecision>
requires (std::is_same_v<TPrecision, float> || std::is_same_v<TPrecision, half>)
class Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >

CUDA implementation of the fused softmax and cross entropy operation for neural networks.

This class provides a CUDA-based implementation of the fused softmax and cross entropy operation, which combines two commonly used operations in neural networks to improve computational efficiency. First, the softmax function converts a vector of real numbers (logits) into a probability distribution. Then, the cross entropy computes the negative log likelihood of the correct class given the predicted probabilities.

The implementation is optimized for NVIDIA GPUs using CUDA for high-performance computation, especially for large vocabulary sizes typical in language models.

Template Parameters

TPrecision The data type used for computation (float or half).

Member Typedef Documentation

◆ MR

template<typename TPrecision >

using Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >::MR = typename CudaDevice::MR

Constructor & Destructor Documentation

◆ FusedSoftmaxCrossEntropyOp() [1/2]

template<typename TPrecision >

Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >::FusedSoftmaxCrossEntropyOp ( )

inline

Constructs a new CUDA Fused Softmax Cross Entropy operation with the default device context.

Initializes the operation with a CUDA device context (defaults to CUDA:0).

◆ FusedSoftmaxCrossEntropyOp() [2/2]

template<typename TPrecision >

Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >::FusedSoftmaxCrossEntropyOp ( std::shared_ptr< DeviceContext > context )

inline

Constructs a new CUDA Fused Softmax Cross Entropy operation with a specific device context.

Parameters

context The device context to use for this operation.

Exceptions

std::runtime_error If the context is not for a CUDA device.

Member Function Documentation

◆ backward()

template<typename TPrecision >

void Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >::backward	(	const Tensor< TPrecision, MR > &	input1,
		const Tensor< int, MR > &	input2,
		const Tensor< TPrecision, MR > &	output,
		const Tensor< TPrecision, MR > &	output_gradient,
		const std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &	parameters,
		std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &	parameter_gradients,
		Tensor< TPrecision, MR > &	input1_gradient,
		Tensor< int, MR > &	input2_gradient,
		const OperationAttributes &	properties,
		const std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &	output_state
	)		const

inlineoverride

Performs the backward pass of the fused softmax cross entropy operation.

Computes gradients with respect to the logits by combining the gradients from both the softmax and cross entropy operations.

Parameters

input1	Input tensor from the forward pass (logits).
input2	Input tensor from the forward pass (target indices).
output	Output tensor from the forward pass (loss values).
output_gradient	Gradient of the loss with respect to the output.
parameters	Parameters tensor from forward pass (not used).
parameter_gradients	Gradients for parameters (not used).
input1_gradient	Gradient of the loss with respect to the logits.
input2_gradient	Not used as we don't compute gradients for integer indices.
properties	Additional attributes for the operation.
output_state	Cache tensors from forward pass containing computed probabilities.

Here is the call graph for this function:

◆ forward()

template<typename TPrecision >

void Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >::forward	(	const Tensor< TPrecision, MR > &	logits,
		const Tensor< int, MR > &	targets,
		const std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &	parameters,
		const OperationAttributes &	properties,
		Tensor< TPrecision, MR > &	losses,
		std::vector< std::shared_ptr< Tensor< TPrecision, MR > > > &	output_state
	)		const

inlineoverride

Performs the forward pass of the fused softmax cross entropy operation on CUDA.

Converts input logits into a probability distribution using softmax and then computes the cross entropy loss between the predicted probabilities and target classes.

Parameters

logits	Input tensor containing logits of shape [B, S, V], where B is batch size, S is sequence length, and V is vocabulary size.
input2	Input tensor containing target class indices of shape [B, S].
parameters	Additional parameters (not used in this operation).
properties	Additional attributes for the operation.
output	Output tensor of shape [B, S] to store the resulting cross entropy loss values.
output_state	Cache for intermediate results, stores the computed probabilities for the backward pass.

Here is the call graph for this function:

◆ getName()

template<typename TPrecision >

std::string Mila::Dnn::Compute::FusedSoftmaxCrossEntropyOp< TPrecision >::getName ( ) const

inlineoverridevirtual

Gets the name of this operation.

Returns: std::string The name of the operation ("Cuda::FusedSoftmaxCrossEntropyOp").

Implements Mila::Dnn::Compute::OperationBase< TDeviceType, TInput1, TInput2, TOutput >.

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

/home/runner/work/Mila/Mila/Mila/Src/Dnn/Compute/Operations/Cuda/FusedSoftmaxCrossEntropyOp.ixx

Public Types

Public Member Functions

Detailed Description

Member Typedef Documentation

◆ MR

Constructor & Destructor Documentation

◆ FusedSoftmaxCrossEntropyOp() [1/2]

◆ FusedSoftmaxCrossEntropyOp() [2/2]

Member Function Documentation

◆ backward()

◆ forward()

◆ getName()