File size: 3,954 Bytes
533920b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import os
from dataclasses import dataclass, field
from typing import Any
import torch
from accelerate.utils import PrecisionType
from accelerate import Accelerator, DeepSpeedPlugin
from omegaconf import OmegaConf, MISSING, II
from trainer.accelerators.base_accelerator import BaseAcceleratorConfig, BaseAccelerator
@dataclass
class MixedPrecisionConfig:
enabled: bool = MISSING
@dataclass
class DeepSpeedConfig:
fp16: MixedPrecisionConfig = field(default_factory=lambda: MixedPrecisionConfig(enabled=False))
bf16: MixedPrecisionConfig = field(default_factory=lambda: MixedPrecisionConfig(enabled=False))
optimizer: dict = field(default_factory=lambda: {
"type": "AdamW",
"params": {
"lr": "auto",
"weight_decay": "auto",
"torch_adam": True,
"adam_w_mode": True
}
})
scheduler: dict = field(default_factory=lambda: {
"type": "WarmupDecayLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto",
"total_num_steps": "auto"
}
})
zero_optimization: dict = field(default_factory=lambda: {
"stage": 2,
"allgather_partitions": True,
"allgather_bucket_size": 2e8,
"overlap_comm": True,
"reduce_scatter": True,
"reduce_bucket_size": 500000000,
"contiguous_gradients": True
})
gradient_accumulation_steps: int = 4
gradient_clipping: float = 1.0
steps_per_print: int = 1
train_batch_size: str = "auto"
train_micro_batch_size_per_gpu: str = "auto"
# train_micro_batch_size_per_gpu: int = II("dataset.batch_size")
wall_clock_breakdown: bool = False
@dataclass
class DeepSpeedAcceleratorConfig(BaseAcceleratorConfig):
_target_: str = "trainer.accelerators.deepspeed_accelerator.DeepSpeedAccelerator"
deepspeed: DeepSpeedConfig = field(default_factory=DeepSpeedConfig)
deepspeed_final: Any = None
class DeepSpeedAccelerator(BaseAccelerator):
def __init__(self, cfg: DeepSpeedAcceleratorConfig):
super().__init__(cfg)
self.set_mixed_precision()
deepspeed_plugin = DeepSpeedPlugin(
hf_ds_config=OmegaConf.to_container(self.cfg.deepspeed, resolve=True),
gradient_accumulation_steps=self.cfg.gradient_accumulation_steps,
)
self.cfg.deepspeed_final = OmegaConf.create(deepspeed_plugin.deepspeed_config)
self.accelerator = Accelerator(
deepspeed_plugin=deepspeed_plugin,
gradient_accumulation_steps=self.cfg.gradient_accumulation_steps,
mixed_precision=self.cfg.mixed_precision,
log_with=self.cfg.log_with,
project_dir=self.cfg.output_dir,
dynamo_backend=self.cfg.dynamo_backend,
)
self.post_init()
def set_mixed_precision(self):
if self.cfg.mixed_precision == PrecisionType.BF16:
self.cfg.deepspeed.bf16.enabled = True
self.cfg.deepspeed.fp16.enabled = False
elif self.cfg.mixed_precision == PrecisionType.FP16:
self.cfg.deepspeed.fp16.enabled = True
self.cfg.deepspeed.bf16.enabled = False
else:
self.cfg.deepspeed.fp16.enabled = False
self.cfg.deepspeed.bf16.enabled = False
def prepare(self, *args, device_placement=None):
prepared = self.accelerator.prepare(*args, device_placement=device_placement)
for obj in prepared:
if isinstance(obj, torch.nn.Module):
if self.cfg.mixed_precision == PrecisionType.BF16:
obj.forward = torch.autocast(device_type=self.device.type, dtype=torch.bfloat16)(obj.forward)
elif self.cfg.mixed_precision == PrecisionType.FP16:
obj.forward = torch.autocast(device_type=self.device.type, dtype=torch.float16)(obj.forward)
return prepared
|