mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-05 19:13:01 +00:00
[npu] change device to accelerator api (#5239)
* update accelerator * fix timer * fix amp * update * fix * update bug * add error raise * fix autocast * fix set device * remove doc accelerator * update doc * update doc * update doc * use nullcontext * update cpu * update null context * change time limit for example * udpate * update * update * update * [npu] polish accelerator code --------- Co-authored-by: Xuanlei Zhao <xuanlei.zhao@gmail.com> Co-authored-by: zxl <43881818+oahzxl@users.noreply.github.com>
This commit is contained in:
@@ -5,6 +5,7 @@ import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from colossalai.accelerator import get_accelerator
|
||||
from colossalai.legacy.communication import (
|
||||
recv_backward,
|
||||
recv_forward,
|
||||
@@ -18,7 +19,6 @@ from colossalai.legacy.core import global_context as gpc
|
||||
from colossalai.legacy.initialize import launch
|
||||
from colossalai.logging import get_dist_logger
|
||||
from colossalai.testing import rerun_if_address_is_in_use, spawn
|
||||
from colossalai.utils import get_current_device
|
||||
|
||||
BATCH_SIZE = 4
|
||||
SEQ_LENGTH = 2
|
||||
@@ -73,7 +73,7 @@ def check_forward_backward(output_tensor, output_grad, rank, logger):
|
||||
|
||||
def check_comm(size, rank, prev_rank, next_rank, logger):
|
||||
dtype = torch.float32
|
||||
device = get_current_device()
|
||||
device = get_accelerator().get_current_device()
|
||||
tensor_shape = (BATCH_SIZE, SEQ_LENGTH, HIDDEN_SIZE)
|
||||
grad_shape = (BATCH_SIZE, SEQ_LENGTH, HIDDEN_SIZE)
|
||||
tensor = torch.randn(tensor_shape, dtype=dtype, device=device)
|
||||
|
Reference in New Issue
Block a user