| |
| |
|
|
| |
|
|
| 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 |
|
|
| |
| 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 |
|
|
