wehub-resource-sync
91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from .optimized_linear import LoRAOptimizedLinear, OptimizedLinear
|
|
|
|
import torch
|
|
|
|
try:
|
|
import transformers
|
|
except ImportError:
|
|
transformers = None
|
|
|
|
|
|
def init_lora(model):
|
|
model.requires_grad_(False)
|
|
for m in model.modules():
|
|
if isinstance(m, LoRAOptimizedLinear):
|
|
m.init_lora()
|
|
|
|
|
|
class Init(object):
|
|
"""
|
|
Init context wrapper similar in style to zero.Init. Allows for injecting OptimizedLinear during model
|
|
construction which will shard base weights and reduce overall memory usage during model init. Primarily
|
|
useful when initializing a model via transformers.AutoModelForCausalLM.
|
|
|
|
Example usage:
|
|
lora_config = deepspeed.linear.LoRAConfig(..)
|
|
quant_config = deepspeed.linear.QuantizationConfig(..)
|
|
with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config):
|
|
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-405B")
|
|
|
|
"""
|
|
|
|
def __init__(self, lora_config=None, quant_config=None):
|
|
self._orig_nn_linear = torch.nn.Linear
|
|
self._orig_causallm_pretrained = None
|
|
if transformers != None:
|
|
self._orig_causallm_pretrained = transformers.AutoModelForCausalLM.from_pretrained
|
|
self._orig_causallm_config = transformers.AutoModelForCausalLM.from_config
|
|
self.lora_config = lora_config
|
|
self.quant_config = quant_config
|
|
self._post_init_complete = False
|
|
|
|
def __enter__(self):
|
|
|
|
class OptLinearWrapper:
|
|
_orig_nn_linear = self._orig_nn_linear
|
|
_lora_config = self.lora_config
|
|
_quant_config = self.quant_config
|
|
|
|
def __new__(self, *args, **kwargs):
|
|
self._lora_config.delay_lora_init = True
|
|
kwargs['lora_config'] = self._lora_config
|
|
kwargs['quantization_config'] = self._quant_config
|
|
kwargs['linear_cls'] = self._orig_nn_linear
|
|
return OptimizedLinear(*args, **kwargs)
|
|
|
|
def _model_init(model):
|
|
if self.lora_config != None:
|
|
init_lora(model)
|
|
self._post_init_complete = True
|
|
return model
|
|
|
|
# ensures non-lora params are frozen and lora weights are initialized
|
|
def from_pretrained(*args, **kwargs):
|
|
model = self._orig_causallm_pretrained(*args, **kwargs)
|
|
return _model_init(model)
|
|
|
|
def from_config(*args, **kwargs):
|
|
model = self._orig_causallm_config(*args, **kwargs)
|
|
return _model_init(model)
|
|
|
|
torch.nn.Linear = OptLinearWrapper
|
|
if transformers != None:
|
|
transformers.AutoModelForCausalLM.from_pretrained = from_pretrained
|
|
transformers.AutoModelForCausalLM.from_config = from_config
|
|
|
|
def __exit__(self, *args, **kwargs):
|
|
torch.nn.Linear = self._orig_nn_linear
|
|
if not self._post_init_complete:
|
|
print('WARNING: For some reason LoRA modules are not initialized, this is usually done automatically '
|
|
'if using transformers via (AutoModelForCausalLM from_pretrained/from_config). '
|
|
'You must call `init_lora` on each module in order to use DeepSpeed LoRA, otherwise '
|
|
'you will error out during runtime.')
|
|
else:
|
|
transformers.AutoModelForCausalLM.from_pretrained = self._orig_causallm_pretrained
|
|
transformers.AutoModelForCausalLM.from_config = self._orig_causallm_config
|