[Feature] auto-cast optimizers to distributed version (#5746)

* auto-cast optimizers to distributed * fix galore casting * logger --------- Co-authored-by: Edenzzzz <wtan45@wisc.edu>
2025-09-05 02:51:59 +00:00 · 2024-05-24 17:24:16 +08:00
parent 2fc85abf43
commit 5f8c0a0ac3
13 changed files with 61 additions and 31 deletions
--- a/colossalai/nn/optimizer/distributed_came.py
+++ b/colossalai/nn/optimizer/distributed_came.py
@@ -34,9 +34,6 @@ class DistributedCAME(DistributedOptim):
        betas=(0.9, 0.999, 0.9999),
        weight_decay=0.0,
    ):
-        assert lr > 0.0
-        assert all([0.0 <= beta <= 1.0 for beta in betas])
-
        defaults = dict(
            lr=lr,
            eps=eps,