66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
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()
|