GitLab

* [sync BN]

support non-uniform batch size across process group.

TODO: test should be added once cleaned up.

* updating unit tests

* new unit tests for different inputs

* cleaning

* Only attempt to patch Tensor methods if defined

* syntax
Co-authored-by: Michael Carilli <mcarilli@nvidia.com>

add unit tests for FusedLAMB optimizer

2d masking and sparsity

Minor edit

Importance and usage is 2d masking

Update softmax.h

Fixing NULL stream launch in dispatch_additive_masked_softmax

add streaming support for softmax kernels

Sparsity

Co-authored-by: Michael Carilli <mcarilli@nvidia.com>

Remove distributed lamb from __init__.py

Distributed lamb optimizer