Spaces:
Running on Zero
Running on Zero
| from email.policy import strict | |
| import os | |
| import copy | |
| from functools import partial | |
| from contextlib import nullcontext | |
| import torch | |
| import torch.distributed as dist | |
| from torch.nn.parallel import DistributedDataParallel as DDP | |
| import numpy as np | |
| from .utils import * | |
| from .base import Trainer | |
| from ..utils.general_utils import * | |
| from ..utils.dist_utils import * | |
| from ..utils import grad_clip_utils, elastic_utils | |
| class BasicTrainer(Trainer): | |
| """ | |
| Trainer for basic training loop. | |
| Args: | |
| models (dict[str, nn.Module]): Models to train. | |
| dataset (torch.utils.data.Dataset): Dataset. | |
| output_dir (str): Output directory. | |
| load_dir (str): Load directory. | |
| step (int): Step to load. | |
| batch_size (int): Batch size. | |
| batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored. | |
| batch_split (int): Split batch with gradient accumulation. | |
| max_steps (int): Max steps. | |
| optimizer (dict): Optimizer config. | |
| lr_scheduler (dict): Learning rate scheduler config. | |
| elastic (dict): Elastic memory management config. | |
| grad_clip (float or dict): Gradient clip config. | |
| ema_rate (float or list): Exponential moving average rates. | |
| fp16_mode (str): FP16 mode. | |
| - None: No FP16. | |
| - 'inflat_all': Hold a inflated fp32 master param for all params. | |
| - 'amp': Automatic mixed precision. | |
| fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation. | |
| finetune_ckpt (dict): Finetune checkpoint. | |
| log_param_stats (bool): Log parameter stats. | |
| i_print (int): Print interval. | |
| i_log (int): Log interval. | |
| i_sample (int): Sample interval. | |
| i_save (int): Save interval. | |
| i_ddpcheck (int): DDP check interval. | |
| """ | |
| def __str__(self): | |
| lines = [] | |
| lines.append(self.__class__.__name__) | |
| lines.append(f' - Models:') | |
| for name, model in self.models.items(): | |
| lines.append(f' - {name}: {model.__class__.__name__}') | |
| lines.append(f' - Dataset: {indent(str(self.dataset), 2)}') | |
| lines.append(f' - Dataloader:') | |
| lines.append(f' - Sampler: {self.dataloader.sampler.__class__.__name__}') | |
| lines.append(f' - Num workers: {self.dataloader.num_workers}') | |
| lines.append(f' - Number of steps: {self.max_steps}') | |
| lines.append(f' - Number of GPUs: {self.world_size}') | |
| lines.append(f' - Batch size: {self.batch_size}') | |
| lines.append(f' - Batch size per GPU: {self.batch_size_per_gpu}') | |
| lines.append(f' - Batch split: {self.batch_split}') | |
| lines.append(f' - Optimizer: {self.optimizer.__class__.__name__}') | |
| lines.append(f' - Learning rate: {self.optimizer.param_groups[0]["lr"]}') | |
| if self.lr_scheduler_config is not None: | |
| lines.append(f' - LR scheduler: {self.lr_scheduler.__class__.__name__}') | |
| if self.elastic_controller_config is not None: | |
| lines.append(f' - Elastic memory: {indent(str(self.elastic_controller), 2)}') | |
| if self.grad_clip is not None: | |
| lines.append(f' - Gradient clip: {indent(str(self.grad_clip), 2)}') | |
| lines.append(f' - EMA rate: {self.ema_rate}') | |
| lines.append(f' - FP16 mode: {self.fp16_mode}') | |
| return '\n'.join(lines) | |
| def init_models_and_more(self, **kwargs): | |
| """ | |
| Initialize models and more. | |
| """ | |
| if self.optimizer_config['name'].lower() == 'muon': | |
| self._ensure_uniform_param_dtype_for_muon() | |
| if self.world_size > 1: | |
| # Prepare distributed data parallel. | |
| # NOTE: PyTorch will raise if a module has no trainable parameters. | |
| self.training_models = {} | |
| for name, model in self.models.items(): | |
| has_trainable_params = any(p.requires_grad for p in model.parameters()) | |
| if has_trainable_params: | |
| self.training_models[name] = DDP( | |
| model, | |
| device_ids=[self.local_rank], | |
| output_device=self.local_rank, | |
| bucket_cap_mb=128, | |
| find_unused_parameters=False | |
| ) | |
| else: | |
| self.training_models[name] = model | |
| if self.is_master: | |
| print(f"\n\033[93mWarning: Model '{name}' has no trainable parameters; skipping DDP wrapping.\033[0m") | |
| else: | |
| self.training_models = self.models | |
| if self.optimizer_config['name'].lower() == 'muon' and self.fp16_mode == 'inflat_all': | |
| self.fp16_mode = 'amp_manual' | |
| print(f'\n\033[93mWarning: Muon optimizer does not support inflat_all mode, switching to amp mode.\033[0m') | |
| # Build master params | |
| # self.model_params = sum([[p for p in model.parameters() if p.requires_grad] for model in self.models.values()], []) | |
| self.model_params_names, self.model_params = [], [] | |
| for key in self.models: | |
| model = self.models[key] | |
| for name, param in model.named_parameters(): | |
| if param.requires_grad: | |
| self.model_params_names.append(key + '.' + name) | |
| self.model_params.append(param) | |
| if 'amp' in self.fp16_mode: | |
| self.master_params = self.model_params | |
| self.scaler = torch.GradScaler() | |
| elif self.fp16_mode == 'inflat_all': | |
| self.master_params = make_master_params(self.model_params) | |
| self.fp16_scale_growth = self.fp16_scale_growth | |
| self.log_scale = 15.0 | |
| elif self.loss_scaling: | |
| self.master_params = make_master_params_regular(self.model_params) | |
| self.log_scale = 15.0 | |
| else: | |
| self.master_params = self.model_params | |
| # Build EMA params | |
| if self.is_master: | |
| self.ema_params = [copy.deepcopy(self.master_params) for _ in self.ema_rate] | |
| # Initialize optimizer | |
| if self.optimizer_config['name'].lower() == 'muon': | |
| from muon import MuonWithAuxAdam | |
| hidden_attn_params, other_params = [], [] | |
| for n, p in zip(self.model_params_names, self.master_params): | |
| attn_only = self.optimizer_config['args'].get('attn_only', False) | |
| if ((attn_only and 'attn' in n) or not attn_only) and p.ndim >= 2: | |
| hidden_attn_params.append(p) | |
| else: | |
| other_params.append(p) | |
| param_groups = [ | |
| dict(use_muon=True, params=hidden_attn_params, lr=self.optimizer_config['args']['muon_lr'], weight_decay=self.optimizer_config['args'].get('muon_weight_decay', 0.01)), | |
| dict(use_muon=False, params=other_params, lr=self.optimizer_config['args']['other_lr'], weight_decay=self.optimizer_config['args'].get('other_weight_decay', 0.01)) | |
| ] | |
| self.optimizer = MuonWithAuxAdam(param_groups) | |
| elif hasattr(torch.optim, self.optimizer_config['name']): | |
| self.optimizer = getattr(torch.optim, self.optimizer_config['name'])(self.master_params, **self.optimizer_config['args']) | |
| else: | |
| self.optimizer = globals()[self.optimizer_config['name']](self.master_params, **self.optimizer_config['args']) | |
| # Initalize learning rate scheduler | |
| if self.lr_scheduler_config is not None: | |
| if hasattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name']): | |
| if self.lr_scheduler_config['name'] == 'SequentialLR': | |
| scheduler_list = [] | |
| for scheduler_config in self.lr_scheduler_config['schedulers']: | |
| scheduler_list.append(getattr(torch.optim.lr_scheduler, scheduler_config['name'])(self.optimizer, **scheduler_config['args'])) | |
| self.lr_scheduler = getattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name'])(self.optimizer, scheduler_list, **self.lr_scheduler_config['args']) | |
| else: | |
| self.lr_scheduler = getattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name'])(self.optimizer, **self.lr_scheduler_config['args']) | |
| else: | |
| self.lr_scheduler = globals()[self.lr_scheduler_config['name']](self.optimizer, **self.lr_scheduler_config['args']) | |
| # Initialize elastic memory controller | |
| if self.elastic_controller_config is not None: | |
| assert any([isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)) for model in self.models.values()]), \ | |
| 'No elastic module found in models, please inherit from ElasticModule or ElasticModuleMixin' | |
| self.elastic_controller = getattr(elastic_utils, self.elastic_controller_config['name'])(**self.elastic_controller_config['args']) | |
| for model in self.models.values(): | |
| if isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)): | |
| model.register_memory_controller(self.elastic_controller) | |
| # Initialize gradient clipper | |
| if self.grad_clip is not None: | |
| if isinstance(self.grad_clip, (float, int)): | |
| self.grad_clip = float(self.grad_clip) | |
| else: | |
| self.grad_clip = getattr(grad_clip_utils, self.grad_clip['name'])(**self.grad_clip['args']) | |
| def _ensure_uniform_param_dtype_for_muon(self): | |
| """Muon optimizer gathers parameters across ranks and requires uniform dtypes.""" | |
| target_dtype = self._determine_muon_param_dtype() | |
| casted_any = False | |
| for model in self.models.values(): | |
| if any(param.dtype != target_dtype for param in model.parameters() if param.requires_grad): | |
| casted_any = True | |
| if not self._cast_model_to_dtype(model, target_dtype): | |
| for param in model.parameters(): | |
| if param.requires_grad and param.dtype != target_dtype: | |
| param.data = param.data.to(dtype=target_dtype) | |
| for buffer in model.buffers(): | |
| if buffer.is_floating_point() and buffer.dtype != target_dtype: | |
| buffer.data = buffer.data.to(dtype=target_dtype) | |
| if casted_any and self.is_master: | |
| print(f"\nMuon optimizer converted trainable parameters to {target_dtype} to keep dtype consistent across ranks.") | |
| def _determine_muon_param_dtype(self): | |
| if isinstance(self.muon_param_dtype, torch.dtype): | |
| return self.muon_param_dtype | |
| if self.muon_param_dtype is None: | |
| return torch.float32 | |
| key = str(self.muon_param_dtype).lower() | |
| if key in ('fp16', 'float16', 'half'): | |
| return torch.float16 | |
| if key in ('bf16', 'bfloat16'): | |
| return torch.bfloat16 | |
| if key in ('fp32', 'float32'): | |
| return torch.float32 | |
| raise ValueError(f'Unsupported muon_param_dtype: {self.muon_param_dtype}') | |
| def _cast_model_to_dtype(self, model, target_dtype): | |
| if target_dtype == torch.float16 and hasattr(model, 'convert_to_fp16'): | |
| model.convert_to_fp16() | |
| return True | |
| if target_dtype == torch.float32 and hasattr(model, 'convert_to_fp32'): | |
| model.convert_to_fp32() | |
| return True | |
| if target_dtype == torch.bfloat16 and hasattr(model, 'convert_to_bf16'): | |
| model.convert_to_bf16() | |
| return True | |
| return False | |
| def _master_params_to_state_dicts(self, master_params): | |
| """ | |
| Convert master params to dict of state_dicts. | |
| """ | |
| if self.fp16_mode == 'inflat_all': | |
| master_params = unflatten_master_params(self.model_params, master_params) | |
| state_dicts = {name: model.state_dict() for name, model in self.models.items()} | |
| master_params_names = sum( | |
| [[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()] | |
| , []) | |
| for i, (model_name, param_name) in enumerate(master_params_names): | |
| state_dicts[model_name][param_name] = master_params[i] | |
| return state_dicts | |
| def _state_dicts_to_master_params(self, master_params, state_dicts, strict=True): | |
| """ | |
| Convert a state_dict to master params. | |
| """ | |
| master_params_names = sum( | |
| [[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()], []) | |
| if strict: | |
| params = [state_dicts[name][param_name] for name, param_name in master_params_names] | |
| else: | |
| param_dict = {key: dict(value.named_parameters()) for key, value in self.models.items()} | |
| params = [] | |
| for i, (name, param_name) in enumerate(master_params_names): | |
| if name in state_dicts and param_name in state_dicts[name]: | |
| params.append(state_dicts[name][param_name].to(master_params[0].device)) | |
| else: | |
| params.append(param_dict[name][param_name].data.clone().to(master_params[0].device)) | |
| if self.fp16_mode == 'inflat_all': | |
| model_params_to_master_params(params, master_params) | |
| else: | |
| for i, param in enumerate(params): | |
| master_params[i].data.copy_(param.data) | |
| def load_pretrain(self, load_dir): | |
| """ | |
| Load pretrained checkpoints. | |
| Should be called by all processes. | |
| """ | |
| if self.is_master: | |
| print(f'\nLoading pretrained checkpoint from step {load_dir}...', end='') | |
| need_load = False | |
| model_ckpts = {} | |
| for name, model in self.models.items(): | |
| if hasattr(model, 'use_pretrain_branch') and model.use_pretrain_branch: | |
| class_names = model.__class__.__name__ if not hasattr(model, 'pretrain_class_name') else model.pretrain_class_name | |
| need_load = True | |
| from safetensors.torch import load_file | |
| if type(class_names) is not list: | |
| class_names = [class_names] | |
| for class_name in class_names: | |
| if os.path.exists(os.path.join('ckpts/pretrained_ckpts', f'{class_name}.safetensors')): | |
| pretrain_ckpt = load_file(os.path.join('ckpts/pretrained_ckpts', f'{class_name}.safetensors')) | |
| elif os.path.exists(os.path.join('ckpts/pretrained_ckpts', f'{class_name}.pt')): | |
| pretrain_ckpt = torch.load(os.path.join('ckpts/pretrained_ckpts', f'{class_name}.pt'), weights_only=True) | |
| else: | |
| continue | |
| if hasattr(model, 'pretrain_ckpt_filter_prefix') and class_name in model.pretrain_ckpt_filter_prefix: | |
| prefix = model.pretrain_ckpt_filter_prefix[class_name] | |
| pretrain_ckpt = {key: value for key, value in pretrain_ckpt.items() if key.startswith(prefix)} | |
| model_ckpts[name] = pretrain_ckpt | |
| model.load_state_dict(pretrain_ckpt, strict=False) | |
| if need_load: | |
| self._state_dicts_to_master_params(self.master_params, model_ckpts, strict=False) | |
| del model_ckpts | |
| if self.world_size > 1: | |
| dist.barrier() | |
| if self.is_master: | |
| print('Load Pretrained Checkpoints Done.') | |
| if self.world_size > 1: | |
| self.check_ddp() | |
| def load(self, load_dir, step=0): | |
| """ | |
| Load a checkpoint. | |
| Should be called by all processes. | |
| """ | |
| if self.is_master: | |
| print(f'\nLoading checkpoint from step {step}...', end='') | |
| model_ckpts = {} | |
| for name, model in self.models.items(): | |
| if os.path.exists(os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt')): | |
| model_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt')), map_location=self.device, weights_only=True) | |
| model_ckpts[name] = model_ckpt | |
| model.load_state_dict(model_ckpt, strict=True) | |
| else: | |
| model_ckpts[name] = model.state_dict() | |
| if self.fp16_mode == 'inflat_all' or self.loss_scaling: | |
| model.convert_to_fp16() | |
| self._state_dicts_to_master_params(self.master_params, model_ckpts, strict=True) | |
| del model_ckpts | |
| if self.is_master: | |
| for i, ema_rate in enumerate(self.ema_rate): | |
| ema_ckpts = {} | |
| for name, model in self.models.items(): | |
| ckpt_path = os.path.join(load_dir, 'ckpts', f'{name}_ema{ema_rate}_step{step:07d}.pt') | |
| ckpt_path = ckpt_path if os.path.exists(ckpt_path) else os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt') # If EMA ckpt not found, load normal ckpt instead | |
| if os.path.exists(ckpt_path): | |
| ema_ckpt = torch.load(ckpt_path, map_location=self.device, weights_only=True) | |
| ema_ckpts[name] = ema_ckpt | |
| self._state_dicts_to_master_params(self.ema_params[i], ema_ckpts, strict=False) | |
| del ema_ckpts | |
| misc_ckpt_path = os.path.join(load_dir, 'ckpts', f'misc_step{step:07d}.pt') | |
| if os.path.exists(misc_ckpt_path): | |
| misc_ckpt = torch.load(read_file_dist(misc_ckpt_path), map_location=torch.device('cpu'), weights_only=False) | |
| self.optimizer.load_state_dict(misc_ckpt['optimizer']) | |
| self.step = misc_ckpt['step'] | |
| self.data_sampler.load_state_dict(misc_ckpt['data_sampler']) | |
| if 'amp' in self.fp16_mode: | |
| self.scaler.load_state_dict(misc_ckpt['scaler']) | |
| elif self.fp16_mode == 'inflat_all' or self.loss_scaling: | |
| self.log_scale = misc_ckpt['log_scale'] | |
| if self.lr_scheduler_config is not None: | |
| self.lr_scheduler.load_state_dict(misc_ckpt['lr_scheduler']) | |
| if self.elastic_controller_config is not None: | |
| self.elastic_controller.load_state_dict(misc_ckpt['elastic_controller']) | |
| if self.grad_clip is not None and not isinstance(self.grad_clip, float): | |
| self.grad_clip.load_state_dict(misc_ckpt['grad_clip']) | |
| del misc_ckpt | |
| else: | |
| self.step = step | |
| if self.world_size > 1: | |
| dist.barrier() | |
| if self.is_master: | |
| print(' Done.') | |
| if self.world_size > 1: | |
| self.check_ddp() | |
| def save(self): | |
| """ | |
| Save a checkpoint. | |
| Should be called only by the rank 0 process. | |
| """ | |
| assert self.is_master, 'save() should be called only by the rank 0 process.' | |
| print(f'\nSaving checkpoint at step {self.step}...', end='') | |
| model_ckpts = self._master_params_to_state_dicts(self.master_params) | |
| for name, model_ckpt in model_ckpts.items(): | |
| torch.save(model_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_step{self.step:07d}.pt')) | |
| for i, ema_rate in enumerate(self.ema_rate): | |
| ema_ckpts = self._master_params_to_state_dicts(self.ema_params[i]) | |
| for name, ema_ckpt in ema_ckpts.items(): | |
| torch.save(ema_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_ema{ema_rate}_step{self.step:07d}.pt')) | |
| misc_ckpt = { | |
| 'optimizer': self.optimizer.state_dict(), | |
| 'step': self.step, | |
| 'data_sampler': self.data_sampler.state_dict(), | |
| } | |
| if 'amp' in self.fp16_mode: | |
| misc_ckpt['scaler'] = self.scaler.state_dict() | |
| elif self.fp16_mode == 'inflat_all' or self.loss_scaling: | |
| misc_ckpt['log_scale'] = self.log_scale | |
| if self.lr_scheduler_config is not None: | |
| misc_ckpt['lr_scheduler'] = self.lr_scheduler.state_dict() | |
| if self.elastic_controller_config is not None: | |
| misc_ckpt['elastic_controller'] = self.elastic_controller.state_dict() | |
| if self.grad_clip is not None and not isinstance(self.grad_clip, float): | |
| misc_ckpt['grad_clip'] = self.grad_clip.state_dict() | |
| torch.save(misc_ckpt, os.path.join(self.output_dir, 'ckpts', f'misc_step{self.step:07d}.pt')) | |
| print(' Done.') | |
| if not hasattr(self, 'model_saved_steps'): | |
| self.model_saved_steps = [self.step] | |
| else: | |
| self.model_saved_steps.append(self.step) | |
| # Deleting old checkpoints and keep the last 5 checkpoints | |
| self.manage_checkpoints() | |
| def manage_checkpoints(self): | |
| def get_step_from_filename(filename, prefix, suffix): | |
| match = re.search(rf'{prefix}(\d+){suffix}', filename) | |
| return int(match.group(1)) if match else None | |
| ckpt_dir = os.path.join(self.output_dir, 'ckpts') | |
| if not os.path.exists(ckpt_dir): | |
| print("Checkpoint directory does not exist, skipping cleanup.") | |
| return | |
| checkpoints = [] | |
| total_steps = [] | |
| for filename in os.listdir(ckpt_dir): | |
| if filename.endswith('.pt'): | |
| step = get_step_from_filename(filename, 'step', '.pt') | |
| if step is not None: | |
| checkpoints.append((step, filename)) | |
| total_steps.append(step) | |
| checkpoints.sort(key=lambda x: x[0]) | |
| total_steps = sorted(np.unique(total_steps).tolist()) | |
| # Avoid deleting steps mod 100000 to 0 | |
| total_steps = [step for step in total_steps if step % 100000 != 0] | |
| # Keep only the latest 2 checkpoints | |
| if len(total_steps) > 2: | |
| old_steps = total_steps[:-2] | |
| old_checkpoints = [(step, filename) for step, filename in checkpoints if step in old_steps] | |
| for step, filename in old_checkpoints: | |
| os.remove(os.path.join(ckpt_dir, filename)) | |
| print(f'\nDeleted old checkpoint file: {filename}') | |
| for old_step, _ in old_checkpoints: | |
| for name in self.models: | |
| model_ckpt = os.path.join(ckpt_dir, f'{name}_step{old_step:07d}.pt') | |
| if os.path.exists(model_ckpt): | |
| os.remove(model_ckpt) | |
| print(f'Deleted model checkpoint: {model_ckpt}') | |
| for ema_rate in self.ema_rate: | |
| ema_ckpt = os.path.join(ckpt_dir, f'{name}_ema{ema_rate}_step{old_step:07d}.pt') | |
| if os.path.exists(ema_ckpt): | |
| os.remove(ema_ckpt) | |
| print(f'Deleted EMA checkpoint: {ema_ckpt}') | |
| misc_ckpt = os.path.join(ckpt_dir, f'misc_step{old_step:07d}.pt') | |
| if os.path.exists(misc_ckpt): | |
| os.remove(misc_ckpt) | |
| print(f'Deleted misc checkpoint: {misc_ckpt}') | |
| def finetune_from(self, finetune_ckpt): | |
| """ | |
| Finetune from a checkpoint. | |
| Should be called by all processes. | |
| """ | |
| if self.is_master: | |
| print('\nFinetuning from:') | |
| for name, path in finetune_ckpt.items(): | |
| print(f' - {name}: {path}') | |
| model_ckpts = {} | |
| for name, model in self.models.items(): | |
| model_state_dict = model.state_dict() | |
| if name in finetune_ckpt: | |
| model_ckpt = torch.load(read_file_dist(finetune_ckpt[name]), map_location=self.device, weights_only=True) | |
| for k, v in model_ckpt.items(): | |
| if model_ckpt[k].shape != model_state_dict[k].shape: | |
| if self.is_master: | |
| print(f'Warning: {k} shape mismatch, {model_ckpt[k].shape} vs {model_state_dict[k].shape}, skipped.') | |
| model_ckpt[k] = model_state_dict[k] | |
| model_ckpts[name] = model_ckpt | |
| model.load_state_dict(model_ckpt) | |
| if self.fp16_mode == 'inflat_all' or self.loss_scaling: | |
| model.convert_to_fp16() | |
| else: | |
| if self.is_master: | |
| print(f'Warning: {name} not found in finetune_ckpt, skipped.') | |
| model_ckpts[name] = model_state_dict | |
| self._state_dicts_to_master_params(self.master_params, model_ckpts) | |
| del model_ckpts | |
| if self.world_size > 1: | |
| dist.barrier() | |
| if self.is_master: | |
| print('Done.') | |
| if self.world_size > 1: | |
| self.check_ddp() | |
| def update_ema(self): | |
| """ | |
| Update exponential moving average. | |
| Should only be called by the rank 0 process. | |
| """ | |
| assert self.is_master, 'update_ema() should be called only by the rank 0 process.' | |
| for i, ema_rate in enumerate(self.ema_rate): | |
| for master_param, ema_param in zip(self.master_params, self.ema_params[i]): | |
| ema_param.detach().mul_(ema_rate).add_(master_param, alpha=1.0 - ema_rate) | |
| def check_ddp(self): | |
| """ | |
| Check if DDP is working properly. | |
| Should be called by all process. | |
| """ | |
| if self.is_master: | |
| print('\nPerforming DDP check...') | |
| if self.is_master: | |
| print('Checking if parameters are consistent across processes...') | |
| dist.barrier() | |
| try: | |
| for n, p in zip(self.model_params_names, self.master_params): | |
| # split to avoid OOM | |
| for i in range(0, p.numel(), 10000000): | |
| sub_size = min(10000000, p.numel() - i) | |
| sub_p = p.detach().reshape(-1)[i:i+sub_size] | |
| # gather from all processes | |
| sub_p_gather = [torch.empty_like(sub_p) for _ in range(self.world_size)] | |
| dist.all_gather(sub_p_gather, sub_p) | |
| # check if equal | |
| assert all([torch.equal(sub_p, sub_p_gather[i]) for i in range(self.world_size)]), 'parameters are not consistent across processes' | |
| except AssertionError as e: | |
| if self.is_master: | |
| print(f'\n\033[91mError: {e}\033[0m') | |
| print('DDP check failed.') | |
| raise e | |
| dist.barrier() | |
| if self.is_master: | |
| print('Done.') | |
| def run_step(self, data_list): | |
| """ | |
| Run a training step. | |
| """ | |
| step_log = {'loss': {}, 'status': {}} | |
| amp_context = partial(torch.autocast, device_type='cuda') if self.fp16_mode == 'amp' else nullcontext | |
| elastic_controller_context = self.elastic_controller.record if self.elastic_controller_config is not None else nullcontext | |
| # Train | |
| losses = [] | |
| statuses = [] | |
| elastic_controller_logs = [] | |
| zero_grad(self.model_params) | |
| for i, mb_data in enumerate(data_list): | |
| ## sync at the end of each batch split | |
| sync_contexts = [getattr(self.training_models[name], 'no_sync', nullcontext) for name in self.training_models] if i != len(data_list) - 1 and self.world_size > 1 else [nullcontext] | |
| with nested_contexts(*sync_contexts), elastic_controller_context(): | |
| with amp_context(): | |
| loss, status = self.training_losses(**mb_data) | |
| l = loss['loss'] / len(data_list) | |
| ## backward | |
| if 'amp' in self.fp16_mode: | |
| self.scaler.scale(l).backward() | |
| elif self.fp16_mode == 'inflat_all': | |
| scaled_l = l * (2 ** self.log_scale) | |
| scaled_l.backward() | |
| elif self.loss_scaling: | |
| # Mirror manual loss scaling path to keep updates unbiased. | |
| scaled_l = l * (2 ** self.log_scale) | |
| scaled_l.backward() | |
| else: | |
| l.backward() | |
| ## log | |
| losses.append(dict_foreach(loss, lambda x: x.item() if isinstance(x, torch.Tensor) else x)) | |
| statuses.append(dict_foreach(status, lambda x: x.item() if isinstance(x, torch.Tensor) else x)) | |
| if self.elastic_controller_config is not None: | |
| elastic_controller_logs.append(self.elastic_controller.log()) | |
| ## gradient clip | |
| if self.grad_clip is not None: | |
| if 'amp' in self.fp16_mode: | |
| self.scaler.unscale_(self.optimizer) | |
| elif self.fp16_mode == 'inflat_all': | |
| model_grads_to_master_grads(self.model_params, self.master_params) | |
| scale = 2 ** self.log_scale | |
| self.master_params[0].grad.mul_(1.0 / scale if scale > 1e-30 else 1.0) | |
| elif self.loss_scaling: | |
| model_grads_to_master_grads_regular(self.model_params, self.master_params) | |
| scale = 2 ** self.log_scale | |
| inv_scale = 1.0 / scale if scale > 1e-30 else 1.0 | |
| for p in self.master_params: | |
| if p.grad is not None: | |
| p.grad.mul_(inv_scale) | |
| if isinstance(self.grad_clip, float): | |
| grad_norm = torch.nn.utils.clip_grad_norm_(self.master_params, self.grad_clip) | |
| else: | |
| grad_norm = self.grad_clip(self.master_params) | |
| if torch.isfinite(grad_norm): | |
| statuses[-1]['grad_norm'] = grad_norm.item() | |
| # Cache params for NaN recovery | |
| if self.step % 10 == 0 and not any(not p.isfinite().all() for p in self.master_params): | |
| self.previous_step = self.step | |
| self.previous_params = [p.clone().detach() for p in self.master_params] | |
| self.previous_params_model = [p.clone().detach() for p in self.model_params] | |
| # # Check which parameter does not have gradients and print its name | |
| # for name, param in self.training_models.items(): | |
| # for p_name, p in param.named_parameters(): | |
| # if p.requires_grad and p.grad is None: | |
| # print(f'\n\033[93mWarning: {name}.{p_name} does not have gradients at step {self.step}.\033[0m') | |
| # exit(0) | |
| ## step | |
| if 'amp' in self.fp16_mode: | |
| prev_scale = self.scaler.get_scale() | |
| self.scaler.step(self.optimizer) | |
| self.scaler.update() | |
| elif self.fp16_mode == 'inflat_all' or self.loss_scaling: | |
| prev_scale = 2 ** self.log_scale | |
| if not any(not p.grad.isfinite().all() for p in self.master_params if p.grad is not None): | |
| if self.grad_clip is None: | |
| if self.fp16_mode == 'inflat_all': | |
| model_grads_to_master_grads(self.model_params, self.master_params) | |
| scale = 2 ** self.log_scale | |
| self.master_params[0].grad.mul_(1.0 / scale if scale > 1e-30 else 1.0) | |
| else: | |
| model_grads_to_master_grads_regular(self.model_params, self.master_params) | |
| scale = 2 ** self.log_scale | |
| inv_scale = 1.0 / scale if scale > 1e-30 else 1.0 | |
| for p in self.master_params: | |
| if p.grad is not None: | |
| p.grad.mul_(inv_scale) | |
| self.optimizer.step() | |
| if self.fp16_mode == 'inflat_all': | |
| master_params_to_model_params(self.model_params, self.master_params) | |
| else: | |
| master_params_to_model_params_regular(self.model_params, self.master_params) | |
| self.log_scale += self.fp16_scale_growth | |
| else: | |
| self.log_scale -= 1 | |
| if self.log_scale < -100: | |
| self.log_scale = -100 | |
| else: | |
| prev_scale = 1.0 | |
| if not any(not p.grad.isfinite().all() for p in self.model_params): | |
| self.optimizer.step() | |
| else: | |
| print('\n\033[93mWarning: NaN detected in gradients. Skipping update.\033[0m') | |
| ## adjust learning rate | |
| if self.lr_scheduler_config is not None: | |
| statuses[-1]['lr'] = self.lr_scheduler.get_last_lr()[0] | |
| self.lr_scheduler.step() | |
| # If parameters are not finite, revert to previous parameters | |
| if any(not p.isfinite().all() for p in self.master_params): | |
| if hasattr(self, 'previous_params'): | |
| print(f'\n\033[91mError: NaN loss detected in parameters at step {self.step}. Reverting to previous parameters at step {self.previous_step}.\033[0m') | |
| for i, p in enumerate(self.master_params): | |
| p.data.copy_(self.previous_params[i].data) | |
| for i, p in enumerate(self.model_params): | |
| p.data.copy_(self.previous_params_model[i].data) | |
| else: | |
| print(f'\n\033[91mError: NaN loss detected in parameters at step {self.step}. No previous parameters to revert to.\033[0m') | |
| raise RuntimeError('NaN loss detected in parameters and no previous parameters to revert to.') | |
| # Logs | |
| step_log['loss'] = dict_reduce(losses, lambda x: np.mean(x)) | |
| step_log['status'] = dict_reduce(statuses, lambda x: np.mean(x), special_func={'min': lambda x: np.min(x), 'max': lambda x: np.max(x)}) | |
| if self.elastic_controller_config is not None: | |
| step_log['elastic'] = dict_reduce(elastic_controller_logs, lambda x: np.mean(x)) | |
| if self.grad_clip is not None: | |
| step_log['grad_clip'] = self.grad_clip if isinstance(self.grad_clip, float) else self.grad_clip.log() | |
| # Check grad and norm of each param | |
| if self.log_param_stats: | |
| for key in self.models: | |
| param_norms = {} | |
| param_grads = {} | |
| for name, param in self.models[key].named_parameters(): | |
| if param.requires_grad: | |
| param_norms[name] = param.norm().item() | |
| if param.grad is not None and torch.isfinite(param.grad).all(): | |
| param_grads[name] = param.grad.norm().item() / prev_scale | |
| step_log[key+'-param_norms'] = param_norms | |
| step_log[key+'-param_grads'] = param_grads | |
| # Update exponential moving average | |
| if self.is_master: | |
| self.update_ema() | |
| return step_log | |