| import os |
| import time |
|
|
| import pytest |
| import torch |
| import torch.distributed as dist |
| import torch.multiprocessing as mp |
| from torch import nn |
|
|
| from lmdeploy.pytorch.distributed import DefaultContext |
| from lmdeploy.pytorch.nn import ParallelEmbedding |
|
|
|
|
| def parallel_emb(rank: int, world_size: int, vocab_size: int, feat_size: int, padding_idx: int, dtype: torch.dtype, |
| x: torch.Tensor, weight: torch.Tensor, result_queue: mp.Queue): |
| dist.init_process_group('nccl', rank=rank, world_size=world_size) |
| gpu_group = dist.new_group(ranks=list(range(world_size)), backend='nccl') |
|
|
| DefaultContext.attn_tp_group.rank = rank |
| DefaultContext.dist_config.attn_tp = world_size |
| DefaultContext.attn_tp_group.gpu_group = gpu_group |
|
|
| model = ParallelEmbedding(vocab_size=vocab_size, |
| hidden_size=feat_size, |
| padding_idx=padding_idx, |
| dtype=dtype, |
| is_tp=True, |
| device=torch.device(type='cuda', index=rank)) |
|
|
| weight = weight.to(torch.device(type='cuda', index=rank)) |
| model.weight_loader(model.weight, weight) |
|
|
| input = x.to(torch.device(type='cuda', index=rank)) |
|
|
| with torch.inference_mode(): |
| out = model(input) |
|
|
| if rank == 0: |
| result_queue.put(mp.reductions.reduce_tensor(out)) |
|
|
| if dist.is_initialized(): |
| dist.destroy_process_group() |
|
|
|
|
| class TestEmbedding: |
|
|
| @pytest.fixture |
| def vocab_size(self, request): |
| yield request.param |
|
|
| @pytest.fixture |
| def feat_size(self, request): |
| yield request.param |
|
|
| @pytest.fixture |
| def padding_idx(self, request): |
| yield request.param |
|
|
| @pytest.fixture |
| def dtype(self, request): |
| yield request.param |
|
|
| @pytest.fixture |
| def tp(self, request): |
| yield request.param |
|
|
| @pytest.fixture |
| def seqlen(self, request): |
| yield request.param |
|
|
| @pytest.fixture |
| def weight(self, vocab_size, feat_size, dtype): |
| yield torch.rand(vocab_size, feat_size, dtype=dtype) |
|
|
| @pytest.fixture |
| def x(self, seqlen, vocab_size): |
| yield torch.randint(low=0, high=vocab_size, size=(seqlen, ), dtype=torch.int32) |
|
|
| @pytest.fixture |
| def gt(self, x, vocab_size, feat_size, padding_idx, dtype, weight): |
| token_emb = nn.Embedding(vocab_size, |
| feat_size, |
| padding_idx=padding_idx, |
| dtype=dtype, |
| device=torch.device(type='cuda', index=0)) |
| token_emb.weight.data.copy_(weight) |
| token_emb._fill_padding_idx_with_zero() |
| input = x.to(torch.device(type='cuda', index=0)) |
| yield token_emb(input) |
|
|
| @pytest.mark.parametrize('vocab_size', [65576, 65533, 3333], indirect=True) |
| @pytest.mark.parametrize('feat_size', [4096, 768], indirect=True) |
| @pytest.mark.parametrize('padding_idx', [None], indirect=True) |
| @pytest.mark.parametrize('seqlen', [1024, 1011, 128], indirect=True) |
| @pytest.mark.parametrize('tp', [2], indirect=True) |
| @pytest.mark.parametrize('dtype', [torch.bfloat16], indirect=True) |
| def test_embedding(self, vocab_size, feat_size, padding_idx, seqlen, tp, dtype, x, weight, gt): |
| os.environ['MASTER_ADDR'] = 'localhost' |
| os.environ['MASTER_PORT'] = '29500' |
| os.environ['NCCL_SOCKET_IFNAME'] = 'lo' |
|
|
| world_size = tp |
| processes = [] |
| mp.set_start_method('spawn', force=True) |
| result_queue = mp.Queue() |
|
|
| for rank in range(world_size): |
| p = mp.Process(target=parallel_emb, |
| args=(rank, world_size, vocab_size, feat_size, padding_idx, dtype, x, weight, result_queue)) |
| processes.append(p) |
| p.start() |
| time.sleep(0.5) |
|
|
| func, args = result_queue.get() |
| out = func(*args) |
|
|
| for p in processes: |
| p.join(timeout=10) |
| if p.is_alive(): |
| p.terminate() |
| p.join(timeout=5) |
| if p.is_alive(): |
| p.kill() |
|
|
| torch.testing.assert_close(out, gt) |
|
|
