wehub-resource-sync
31 lines
831 B
Python
31 lines
831 B
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
|
|
from transformers import AutoModelForCausalLM
|
|
import deepspeed
|
|
import argparse
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
deepspeed.runtime.utils.see_memory_usage('pre test', force=True)
|
|
|
|
model = AutoModelForCausalLM.from_pretrained('facebook/opt-350M').half().to(get_accelerator().device_name())
|
|
parser = argparse.ArgumentParser()
|
|
parser = deepspeed.add_config_arguments(parser)
|
|
args = parser.parse_args()
|
|
|
|
deepspeed.runtime.utils.see_memory_usage('post test', force=True)
|
|
|
|
m, _, _, _ = deepspeed.initialize(model=model, args=args, enable_hybrid_engine=True)
|
|
|
|
m.eval()
|
|
input = torch.ones(1, 16, device='cuda', dtype=torch.long)
|
|
out = m(input)
|
|
|
|
m.train()
|
|
out = m(input)
|
|
print(out['logits'], out['logits'].norm())
|