anigen/trainers/basic.py

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