ColossalAI/tests/kit/model_zoo/transformers/opt.py

86 lines
3.2 KiB
Python
Raw Normal View History

import torch
import transformers
from ..registry import ModelAttribute, model_zoo
# ===============================
# Register single-sentence OPT
# ===============================
BATCH_SIZE = 2
SEQ_LENGTH = 16
def data_gen():
input_ids = torch.Tensor([[1, 15043, 29892, 590, 11203, 338, 274, 1082]]).long()
attention_mask = torch.Tensor([[1, 1, 1, 1, 1, 1, 1, 1]]).long()
return dict(input_ids=input_ids, attention_mask=attention_mask)
def data_gen_for_causal_lm():
# LM data gen
# the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels`
data = data_gen()
labels = data['input_ids'].clone()
data['labels'] = labels
return data
def data_gen_for_sequence_classification():
# LM data gen
# the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels`
data = data_gen()
labels = data['input_ids'].clone()
data['labels'] = torch.tensor([1])
return data
def data_gen_for_question_answering():
# LM data gen
# the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels`
data = data_gen()
data['start_positions'] = torch.tensor([0])
data['end_positions'] = torch.tensor([1])
return data
output_transform_fn = lambda x: x
loss_fn_for_opt_model = lambda x: torch.nn.functional.mse_loss(x.last_hidden_state, torch.ones_like(x.last_hidden_state)
)
loss_fn_for_lm = lambda x: x.loss
config = transformers.OPTConfig(
hidden_size=128,
num_hidden_layers=2,
num_attention_heads=4,
dropout=0,
)
# register the following models
# transformers.OPTModel,
# transformers.OPTForCausalLM,
model_zoo.register(name='transformers_opt',
model_fn=lambda: transformers.OPTModel(config),
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_opt_model,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_opt_for_causal_lm',
model_fn=lambda: transformers.OPTForCausalLM(config),
data_gen_fn=data_gen_for_causal_lm,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_lm,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name='transformers_opt_for_question_answering',
model_fn=lambda: transformers.OPTForQuestionAnswering(config),
data_gen_fn=data_gen_for_question_answering,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_lm,
model_attribute=ModelAttribute(has_control_flow=True))
# TODO The loss and gradient check in the test are failing, to be fixed.
# model_zoo.register(name='transformers_opt_for_sequence_classification',
# model_fn=lambda: transformers.OPTForSequenceClassification(config),
# data_gen_fn=data_gen_for_sequence_classification,
# output_transform_fn=output_transform_fn,
# loss_fn=loss_fn_for_lm,
# model_attribute=ModelAttribute(has_control_flow=True))