mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-01-29 21:49:54 +00:00
890774b2fb
* [shardformer] support lazy init * [shardformer] linear support lazy init * [shardformer] embedding support lazy init * [shardformer] norm support lazy init * [shardformer] fused linear support lazy init * [test] update shardformer test layer * [test] shardformer with lazy init fit ddp * [lazy] hotfix deepcopy of param * [shardformer] fix bert policy and update test * [shardformer] fix bloom policy and update test * [shardformer] fix opt policy and update test * [shardformer] fix t5 policy and update test * [shardformer] fix gpt2 policy and update test * [shardformer] fix llama policy and update test
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
from contextlib import nullcontext
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from torch.testing import assert_close
|
|
|
|
import colossalai
|
|
from colossalai.lazy import LazyInitContext
|
|
from colossalai.shardformer.layer import Embedding1D
|
|
from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
|
|
|
|
|
|
@parameterize('lazy_init', [False, True])
|
|
def check_embedding_1d(lazy_init: bool):
|
|
ctx = LazyInitContext() if lazy_init else nullcontext()
|
|
|
|
with ctx:
|
|
embedding = nn.Embedding(32, 128).cuda()
|
|
embedding_1d = Embedding1D.from_native_module(embedding, process_group=None)
|
|
|
|
assert embedding_1d.weight.shape == torch.Size([32, 64])
|
|
|
|
# ensure state dict is reversibly loadable
|
|
embedding.load_state_dict(embedding_1d.state_dict())
|
|
embedding_1d.load_state_dict(embedding.state_dict())
|
|
|
|
# check computation correctness
|
|
x = torch.randint(low=0, high=32, size=(4, 32)).cuda()
|
|
out = embedding(x)
|
|
gather_out = embedding_1d(x)
|
|
assert_close(out, gather_out)
|
|
|
|
# check backward correctness
|
|
out.sum().backward()
|
|
gather_out.sum().backward()
|
|
|
|
rank = dist.get_rank()
|
|
target_grad = torch.chunk(embedding.weight.grad, 2, dim=1)[rank]
|
|
assert_close(target_grad, embedding_1d.weight.grad)
|
|
|
|
|
|
def run_dist(rank, world_size, port):
|
|
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
|
|
check_embedding_1d()
|
|
|
|
|
|
@rerun_if_address_is_in_use()
|
|
def test_embedding_1d():
|
|
spawn(run_dist, nprocs=2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_embedding_1d()
|