[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:
Hongxin Liu
2025-02-11 16:10:25 +08:00
committed by GitHub
parent 17062c83b9
commit 2b415e5999
13 changed files with 612 additions and 22 deletions

View File

@@ -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()