mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-07 12:01:39 +00:00
[shardformer] support ep for deepseek v3 (#6185)
* [feature] support ep for deepseek v3 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix test * [shardformer] fix deepseek v3 init * [lazy] fit lora for lazy init * [example] support npu for deepseek v3 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -7,6 +7,7 @@ from torch import Tensor
|
||||
from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler
|
||||
|
||||
from colossalai.cluster import DistCoordinator
|
||||
from colossalai.utils import get_current_device
|
||||
|
||||
|
||||
def divide(x: float, y: float) -> float:
|
||||
@@ -29,7 +30,7 @@ def all_reduce_mean(x: float, world_size: int) -> float:
|
||||
# tensor = tensor / world_size
|
||||
# return tensor.item()
|
||||
|
||||
tensor = torch.tensor([x], device=torch.cuda.current_device(), dtype=torch.float)
|
||||
tensor = torch.tensor([x], device=get_current_device(), dtype=torch.float)
|
||||
dist.all_reduce(tensor)
|
||||
tensor = tensor / world_size
|
||||
return tensor.item()
|
||||
|
Reference in New Issue
Block a user