Make Nora+NAdamW the default optimizer

2026-07-31 23:16:39 +08:00
parent 7aa5ed09d9
commit 04899a2b15
11 changed files with 1010 additions and 105 deletions
@@ -0,0 +1,65 @@
+import pytest
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import fully_shard
+from torch.distributed.tensor import DTensor, Shard
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from astrai.model import AutoRegressiveLM
+from astrai.optim import NoraNAdamW
+from astrai.parallel.setup import find_free_port
+from tests.helpers import make_tiny_config
+
+pytestmark = pytest.mark.skipif(
+    torch.cuda.device_count() < 1, reason="CUDA device required"
+)
+
+
+def _assign_grads_and_step(model):
+    optimizer = NoraNAdamW(model)
+    for param in model.parameters():
+        if param.requires_grad:
+            param.grad = torch.ones_like(param)
+    optimizer.step()
+    return optimizer
+
+
+def test_nora_nadamw_steps_after_ddp_and_fsdp2_wrapping():
+    torch.cuda.set_device(0)
+    dist.init_process_group(
+        "nccl",
+        rank=0,
+        world_size=1,
+        init_method=f"tcp://127.0.0.1:{find_free_port()}",
+    )
+    try:
+        ddp_model = AutoRegressiveLM(make_tiny_config()).to(
+            device="cuda", dtype=torch.bfloat16
+        )
+        ddp_model = DDP(ddp_model, device_ids=[0], output_device=0)
+        ddp_optimizer = _assign_grads_and_step(ddp_model)
+        assert ddp_optimizer.state_dict()["nora"]["state"]
+
+        fsdp_model = AutoRegressiveLM(make_tiny_config()).to(
+            device="cuda", dtype=torch.bfloat16
+        )
+        for child in fsdp_model.children():
+            if isinstance(child, torch.nn.ModuleList):
+                for submodule in child:
+                    fully_shard(submodule, reshard_after_forward=False)
+            else:
+                fully_shard(child, reshard_after_forward=False)
+
+        fsdp_optimizer = _assign_grads_and_step(fsdp_model)
+        nora_params = fsdp_optimizer.nora.param_groups[0]["params"]
+        assert nora_params
+        assert all(isinstance(param, DTensor) for param in nora_params)
+        assert all(
+            all(
+                not isinstance(placement, Shard) or placement.dim == 0
+                for placement in param.placements
+            )
+            for param in nora_params
+        )
+    finally:
+        dist.destroy_process_group()