mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-30 22:24:21 +00:00
[legacy] clean up legacy code (#4743)
* [legacy] remove outdated codes of pipeline (#4692) * [legacy] remove cli of benchmark and update optim (#4690) * [legacy] remove cli of benchmark and update optim * [doc] fix cli doc test * [legacy] fix engine clip grad norm * [legacy] remove outdated colo tensor (#4694) * [legacy] remove outdated colo tensor * [test] fix test import * [legacy] move outdated zero to legacy (#4696) * [legacy] clean up utils (#4700) * [legacy] clean up utils * [example] update examples * [legacy] clean up amp * [legacy] fix amp module * [legacy] clean up gpc (#4742) * [legacy] clean up context * [legacy] clean core, constants and global vars * [legacy] refactor initialize * [example] fix examples ci * [example] fix examples ci * [legacy] fix tests * [example] fix gpt example * [example] fix examples ci * [devops] fix ci installation * [example] fix examples ci
This commit is contained in:
@@ -1,7 +1,8 @@
|
||||
from colossalai.core import global_context as gpc
|
||||
from colossalai.context import ParallelMode
|
||||
import torch
|
||||
|
||||
from colossalai.legacy.context import ParallelMode
|
||||
from colossalai.legacy.core import global_context as gpc
|
||||
|
||||
_MAX_DATA_DIM = 5
|
||||
|
||||
|
||||
@@ -22,7 +23,8 @@ def _build_key_size_numel_dictionaries(keys, data):
|
||||
|
||||
# Move to GPU and broadcast.
|
||||
sizes_cuda = torch.cuda.LongTensor(sizes)
|
||||
torch.distributed.broadcast(sizes_cuda, gpc.get_ranks_in_group(ParallelMode.TENSOR)[0],
|
||||
torch.distributed.broadcast(sizes_cuda,
|
||||
gpc.get_ranks_in_group(ParallelMode.TENSOR)[0],
|
||||
group=gpc.get_group(ParallelMode.TENSOR))
|
||||
|
||||
# Move back to cpu and unpack.
|
||||
@@ -60,19 +62,15 @@ def broadcast_data(keys, data, datatype):
|
||||
"""
|
||||
# Build (key, size) and (key, number of elements) dictionaries along
|
||||
# with the total number of elements on all ranks.
|
||||
key_size, key_numel, total_numel = _build_key_size_numel_dictionaries(keys,
|
||||
data)
|
||||
key_size, key_numel, total_numel = _build_key_size_numel_dictionaries(keys, data)
|
||||
|
||||
# Pack on rank zero.
|
||||
if not gpc.is_initialized(ParallelMode.TENSOR) or gpc.get_local_rank(ParallelMode.TENSOR) == 0:
|
||||
# Check that all keys have the same data type.
|
||||
# Flatten the data associated with the keys
|
||||
flatten_data = torch.cat(
|
||||
[data[key].contiguous().view(-1) for key in keys], dim=0).cuda()
|
||||
flatten_data = torch.cat([data[key].contiguous().view(-1) for key in keys], dim=0).cuda()
|
||||
else:
|
||||
flatten_data = torch.empty(total_numel,
|
||||
device=torch.cuda.current_device(),
|
||||
dtype=datatype)
|
||||
flatten_data = torch.empty(total_numel, device=torch.cuda.current_device(), dtype=datatype)
|
||||
|
||||
# Broadcast
|
||||
torch.distributed.broadcast(flatten_data,
|
||||
@@ -139,7 +137,7 @@ def get_batch_for_sequence_parallel(data_iterator):
|
||||
seq_length = data_b['text'].size(1)
|
||||
sub_seq_length = seq_length // local_world_size
|
||||
sub_seq_start = local_rank * sub_seq_length
|
||||
sub_seq_end = (local_rank+1) * sub_seq_length
|
||||
sub_seq_end = (local_rank + 1) * sub_seq_length
|
||||
#
|
||||
# # Unpack.
|
||||
tokens = data_b['text'][:, sub_seq_start:sub_seq_end].long()
|
||||
@@ -156,10 +154,9 @@ class SequenceParallelDataIterator:
|
||||
|
||||
def __init__(self, data_iter):
|
||||
self.data_iter = data_iter
|
||||
|
||||
|
||||
def __iter__(self):
|
||||
return self.data_iter
|
||||
|
||||
def __next__(self):
|
||||
return get_batch_for_sequence_parallel(self.data_iter)
|
||||
return get_batch_for_sequence_parallel(self.data_iter)
|
||||
|
Reference in New Issue
Block a user