2022-05-07 07:49:14 +00:00
|
|
|
from colossalai.tensor.spec import ShardPattern
|
2022-04-25 08:01:52 +00:00
|
|
|
import torch
|
|
|
|
from colossalai.tensor.op_wrapper import colo_op_impl
|
|
|
|
from colossalai.tensor import ColoTensor
|
2022-05-07 07:49:14 +00:00
|
|
|
from colossalai.nn.loss.loss_1d import VocabParallelCrossEntropyLoss1D
|
2022-04-25 08:01:52 +00:00
|
|
|
|
|
|
|
@colo_op_impl(torch.nn.functional.cross_entropy)
|
|
|
|
def colo_cross_entropy(types, args=(), kwargs=None, pg=None):
|
|
|
|
arg_num = len(args)
|
|
|
|
|
|
|
|
if arg_num > 0:
|
|
|
|
input_tensor = args[0]
|
|
|
|
if arg_num > 1:
|
|
|
|
target = args[1]
|
|
|
|
if arg_num > 2:
|
2022-05-07 07:49:14 +00:00
|
|
|
weight = args[2]
|
2022-04-25 08:01:52 +00:00
|
|
|
|
|
|
|
if 'input' in kwargs:
|
2022-05-07 07:49:14 +00:00
|
|
|
input_tensor = kwargs.pop('input')
|
2022-04-25 08:01:52 +00:00
|
|
|
if 'target' in kwargs:
|
2022-05-07 07:49:14 +00:00
|
|
|
target = kwargs.pop('target')
|
2022-04-25 08:01:52 +00:00
|
|
|
if 'weight' in kwargs:
|
2022-05-07 07:49:14 +00:00
|
|
|
weight = kwargs.pop('weight')
|
2022-04-25 08:01:52 +00:00
|
|
|
|
2022-05-07 07:49:14 +00:00
|
|
|
if not isinstance(input_tensor, ColoTensor):
|
|
|
|
input_tensor = ColoTensor.init_from_torch_tensor(input_tensor)
|
2022-04-25 08:01:52 +00:00
|
|
|
if isinstance(target, ColoTensor):
|
|
|
|
target = target.torch_tensor()
|
|
|
|
|
2022-05-07 07:49:14 +00:00
|
|
|
if input_tensor.is_gathered(): # Input is gathered
|
|
|
|
# TODO(jzy) Shall we make the result of loss function a ColoTensor?
|
|
|
|
return ColoTensor.init_from_torch_tensor(torch.nn.functional.cross_entropy(
|
|
|
|
input_tensor.torch_tensor(), target, weight))
|
|
|
|
elif input_tensor.has_spec() and input_tensor.shard_spec.num_action == 1: # Single Model Parallel Applied
|
|
|
|
if input_tensor.shard_pattern == ShardPattern.Col:
|
|
|
|
return ColoTensor.init_from_torch_tensor(
|
|
|
|
VocabParallelCrossEntropyLoss1D()(input_tensor.torch_tensor(), target))
|
|
|
|
else:
|
|
|
|
raise NotImplementedError
|
|
|
|
else:
|
|
|
|
raise NotImplementedError
|