github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-18 16:00:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

[gemini] fixes for benchmarking (#5847)

Browse Source 
* [gemini] fix missing return

* [gemini] fix missing arg pass

* [gemini] use gather tensor instead of list

* [test] enable flash attention for benchmark by default

* [test] enable flash attention for benchmark by default

---------

Co-authored-by: genghaozhe <939857490@qq.com>
This commit is contained in:
botbw
2024-06-26 15:52:09 +08:00
committed by GitHub
parent 2a25a2aff7
commit 8e718a1421
5 changed files with 27 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/booster/plugin/gemini_plugin.py
View File

@@ -369,9 +369,9 @@ class GeminiPlugin(DPPluginBase):
assert precision in SUPPORTED_PRECISION, f"precision {precision} is not supported"
if get_accelerator().name == "npu":
assert placement_policy == "static", "NPU only supports static placement policy"
if placement_policy == "auto" and enable_async_reduce:
if enable_async_reduce and not pin_memory:
logging.warning(
f"enable_async_reduce requires pin_memory to achieve best performance, which is not implicitly set."
f"enable_async_reduce sets pin_memory=True to achieve best performance, which is not implicitly set."
)
pin_memory = True
self.gemini_config = dict(

12
colossalai/zero/gemini/chunk/chunk.py
View File

@@ -403,9 +403,9 @@ class Chunk:
self.shard_size, dtype=self.dtype, device=get_accelerator().get_current_device()
)
input_list = list(torch.chunk(self.cuda_global_chunk, chunks=self.pg_size, dim=0))
self.grad_reduce_work = dist.reduce_scatter(
self.cuda_shard, input_list, group=self.torch_pg, async_op=async_op
assert self.cuda_global_chunk.is_contiguous()
self.grad_reduce_work = dist.reduce_scatter_tensor(
self.cuda_shard, self.cuda_global_chunk, group=self.torch_pg, async_op=async_op
)
if self.extra_dp_group is not None:
@@ -520,8 +520,10 @@ class Chunk:
assert self.cuda_shard is not None
alloc_storage(self.cuda_global_chunk)
gather_list = list(torch.chunk(input=self.cuda_global_chunk, chunks=self.pg_size, dim=0))
work = dist.all_gather(gather_list, self.cuda_shard, self.torch_pg, async_op=async_op)
assert self.cuda_global_chunk.is_contiguous()
work = dist.all_gather_into_tensor(
self.cuda_global_chunk, self.cuda_shard, self.torch_pg, async_op=async_op
)
self.cuda_shard = None
self.is_gathered = True

4
colossalai/zero/gemini/chunk/manager.py
View File

@@ -133,12 +133,12 @@ class ChunkManager:
self.__sub_accessed_chunk(chunk)
self.__add_memory_usage(chunk.memory_usage)
def move_chunk(self, chunk: Chunk, device: torch.device, force_copy: bool = False) -> None:
def move_chunk(self, chunk: Chunk, device: torch.device, force_copy: bool = False, async_move=False) -> None:
"""Move the shard of the chunk to the target device."""
if not chunk.can_move or chunk.device_type == device.type:
return
self.__sub_memory_usage(chunk.memory_usage)
chunk.shard_move(device, force_copy)
chunk.shard_move(device, force_copy, non_blocking=async_move)
self.__add_memory_usage(chunk.memory_usage)
def trans_tensor_state(self, tensor: torch.Tensor, state: TensorState) -> None:

11
colossalai/zero/gemini/gemini_ddp.py
View File

@@ -387,6 +387,7 @@ class GeminiDDP(ModelWrapper):
p: nn.Parameter,
async_reduce_stream: Optional[torch.cuda.Stream] = None,
):
async_reduce_scatter = async_reduce_stream is not None
setattr(p, "_gemini_reduced", True)
empty_grad = torch.empty_like(grad)
free_storage(empty_grad)
@@ -426,7 +427,7 @@ class GeminiDDP(ModelWrapper):
async_reduce_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(async_reduce_stream):
reduced = chunk_manager.reduce_chunk(grad_chunk, async_op=(async_reduce_stream is not None))
reduced = chunk_manager.reduce_chunk(grad_chunk, async_op=async_reduce_scatter)
if reduced:
grad_chunk.wait_async_reduce()
if not chunk_manager.reuse_fp16_chunk:
@@ -447,9 +448,13 @@ class GeminiDDP(ModelWrapper):
# record l2 norm for gradient clipping. flag is bound to fp16 chunk
if chunk.l2_norm_flag:
grad_chunk.set_l2_norm()
chunk_manager.move_chunk(grad_chunk, grads_device[p], force_copy=True)
chunk_manager.move_chunk(
grad_chunk, grads_device[p], force_copy=True, async_move=async_reduce_scatter
)
if not (master_weights) or (enable_gradient_accumulation):
chunk_manager.move_chunk(chunk, grads_device[p], force_copy=True)
chunk_manager.move_chunk(
chunk, grads_device[p], force_copy=True, async_move=async_reduce_scatter
)
return empty_grad
def zero_grad(self, set_to_none: bool = False) -> None: