GitLab

* Adding Amax Primitive and related args.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Enable local-amax for current-scaling and optionally run AR aross FSDP/TP/SP.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding doc for Amax Primitive.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the function name conflict.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Modification as feedback suggested.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix errors from lint.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the wrong amax-scope in the bwd.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Added more description for amax-scope
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the wrong attribute name.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Keep dim for AmaxCalcuation.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Remove keepDim and add shardy_rule
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix shardy_rule
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Remove extra-collective bytes from ref_coll_count due to local amax.
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>