utils.py

import math
import os
from typing import Dict, List

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import transformers
import webdataset as wds
from deepspeed import zero
from transformers import Trainer
from transformers.trainer import is_sagemaker_mp_enabled
from transformers.utils import logging

logger = logging.get_logger(__name__)

if is_sagemaker_mp_enabled():
    import smdistributed.modelparallel.torch as smp

    # from transformers.trainer_pt_utils import smp_forward_backward, smp_forward_only, smp_gather, smp_nested_concat


def find_all_linear_names(named_modules: Dict, target_modules: List[str]):
    cls = torch.nn.Linear
    lora_module_names = set()
    for name, module in named_modules.items():
        if not any([module_name in name for module_name in target_modules]):
            continue

        if isinstance(module, cls):
            lora_module_names.add(name)

    for name in list(lora_module_names):
        if 'lm_head' in name:  # needed for 16-bit
            lora_module_names.remove(name)

    return list(lora_module_names)


def rank0_print(*args):
    if dist.is_initialized():
        if dist.get_rank() == 0:
            print(*args)


def maybe_zero_3(param):
    if hasattr(param, "ds_id"):
        with zero.GatheredParameters([param]):
            param = param.data.detach().cpu().clone()
    else:
        param = param.detach().cpu().clone()
    return param


# Borrowed from peft.utils.get_peft_model_state_dict
def get_peft_state_maybe_zero_3(named_params, bias):
    if bias == "none":
        to_return = {k: t for k, t in named_params if "lora_" in k}
    elif bias == "all":
        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
    elif bias == "lora_only":
        to_return = {}
        maybe_lora_bias = {}
        lora_bias_names = set()
        for k, t in named_params:
            if "lora_" in k:
                to_return[k] = t
                bias_name = k.split("lora_")[0] + "bias"
                lora_bias_names.add(bias_name)
            elif "bias" in k:
                maybe_lora_bias[k] = t
        for k, t in maybe_lora_bias:
            if bias_name in lora_bias_names:
                to_return[bias_name] = t
    else:
        raise NotImplementedError
    to_return = {k: maybe_zero_3(v) for k, v in to_return.items()}
    return to_return


def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
    """Collects the state dict and dump to disk."""
    if trainer.deepspeed:
        torch.cuda.synchronize()
        trainer.save_model(output_dir)
        return

    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {
            key: value.cpu()
            for key, value in state_dict.items()
        }
        del state_dict
        trainer._save(output_dir, state_dict=cpu_state_dict)


class MultiTaskSequentialWdsLoader:
    def __init__(self, loaders, shared_epoch, task_sequence, switch_every_n_batches=1, samples_dict=None):
        self.loaders = loaders
        self.shared_epoch = shared_epoch
        self.task_sequence = task_sequence
        self.switch_every_n_batches = switch_every_n_batches
        self.iters = {}
        self.batch_count = 0
        self.finished_tasks = {task: False for task in task_sequence}
        self.all_finished = False
        self.samples_dict = samples_dict
        self.set_prob()

    def set_prob(self):
        sample_list = []
        available_tasks = []
        for task in self.finished_tasks:
            if not self.finished_tasks[task]:
                sample_list.append(self.samples_dict[task])
                available_tasks.append(task)

        if len(available_tasks) == 0:
            return
        prob_list = [sample / sum(sample_list) for sample in sample_list]
        assert np.isclose(np.array(prob_list).sum(), 1), f'prob must sum to 1, but got {np.array(prob_list).sum()}'
        probs = []
        cum_prob = 0
        for p in prob_list:
            cum_prob += p
            probs.append(cum_prob)

        self.prob_task = {prob: task for prob, task in zip(probs, available_tasks)}

    def current_task(self):
        idx = (self.batch_count // self.switch_every_n_batches) % len(self.task_sequence)

        for offset in range(len(self.task_sequence)):
            candidate_idx = (idx + offset) % len(self.task_sequence)
            candidate_task = self.task_sequence[candidate_idx]
            if not self.finished_tasks[candidate_task]:
                return candidate_task
        # Return None if all tasks are finished
        return None

    def current_task_prob(self):
        if self.all_finished:
            return None
        prob = torch.rand(1)
        for p in self.prob_task:
            if prob.item() <= p:
                return self.prob_task[p]

    def set_epoch(self, epoch):
        for task, shared_epoch in self.shared_epoch.items():
            shared_epoch.set_value(epoch)

    def __iter__(self):
        self.iters = {task: iter(loader) for task, loader in self.loaders.items()}
        self.finished_tasks = {task: False for task in self.task_sequence}
        self.all_finished = all(v for v in self.finished_tasks.values())
        self.set_prob()
        return self

    def __next__(self):
        task = self.current_task()
        # task = self.current_task_prob()
        if task is None:
            raise StopIteration

        try:
            batch = next(self.iters[task])
        except StopIteration:
            self.finished_tasks[task] = True
            self.all_finished = all(v for v in self.finished_tasks.values())
            self.set_prob()  # Reset the task extraction probabilities
            return self.__next__()

        self.batch_count += 1
        return batch


class MultiTaskTrainerWithIterativeResumeWebLoader(Trainer):

    def get_train_dataloader(self):

        round_fn = math.floor

        world_size = int(os.environ['WORLD_SIZE'])
        task_sequence = self.args.task_sequence

        self.loaders = {}
        self.shared_epoch = {}
        samples_dict = {}
        for task in task_sequence:
            dataset, shared_epoch = self.train_dataset[task]['dataset'], self.train_dataset[task]['shared_epoch']
            # task_num_samples = getattr(self.args, f"{task}_num_samples")
            if task == 'depth':
                global_batch_size = self.args.depth_batch_size * world_size
            elif task == 'segmentation':
                global_batch_size = self.args.segmentation_batch_size * world_size
            elif task == 'captioning':
                global_batch_size = self.args.captioning_batch_size * world_size
            # num_batches = round_fn(task_num_samples / global_batch_size)
            num_batches = getattr(self.args, f"{task}_num_batches")
            num_workers = max(1, self.args.dataloader_num_workers)
            num_worker_batches = round_fn(num_batches / num_workers)  # per dataloader worker
            num_batches = num_worker_batches * num_workers
            num_samples = num_batches * global_batch_size
            dataloader = wds.WebLoader(dataset, batch_size=None, shuffle=False, num_workers=self.args.dataloader_num_workers, persistent_workers=self.args.dataloader_num_workers > 0, pin_memory=True)
            # assert task_num_samples == num_samples
            dataloader.num_batches = num_batches
            dataloader.num_samples = num_samples
            self.loaders[task] = dataloader
            self.shared_epoch[task] = shared_epoch
            samples_dict[task] = num_batches

        self.multi_task_loaders = MultiTaskSequentialWdsLoader(self.loaders, self.shared_epoch, task_sequence, switch_every_n_batches=1, samples_dict=samples_dict)

        # return self.multi_task_loaders
        return self.accelerator.prepare(self.multi_task_loaders)

    def create_optimizer(self):
        """
        Setup the optimizer.

        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
        """
        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
        if self.optimizer is None:
            decay_parameters = self.get_decay_parameter_names(opt_model)
            if self.args.vision_lr > 0 and self.args.captioning_projection_lr > 0 and self.args.text_decoder_lr > 0:
                vision_backbone_parameters = [name for name, _ in opt_model.named_parameters() if ('visual_backbone' in name and 'pooling_head' not in name)]
                captioning_projection_parameters = [name for name, _ in opt_model.named_parameters() if 'merge_projection' in name]
                text_decoder_parameters = [name for name, _ in opt_model.named_parameters() if 'text_decoder' in name]
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in vision_backbone_parameters and n not in captioning_projection_parameters and n not in text_decoder_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in vision_backbone_parameters and n not in captioning_projection_parameters and n not in text_decoder_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.vision_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.vision_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in captioning_projection_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.captioning_projection_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in captioning_projection_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.captioning_projection_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in text_decoder_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.text_decoder_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in text_decoder_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.text_decoder_lr,
                    },
                ]
            elif self.args.vision_lr > 0 and self.args.captioning_projection_lr > 0:
                vision_backbone_parameters = [name for name, _ in opt_model.named_parameters() if ('visual_backbone' in name and 'pooling_head' not in name)]
                captioning_projection_parameters = [name for name, _ in opt_model.named_parameters() if 'merge_projection' in name]
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in vision_backbone_parameters and n not in captioning_projection_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in vision_backbone_parameters and n not in captioning_projection_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.vision_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.vision_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in captioning_projection_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.captioning_projection_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in captioning_projection_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.captioning_projection_lr,
                    },
                ]
            # specific lr for vision backbone
            elif self.args.vision_lr > 0:
                vision_backbone_parameters = [name for name, _ in opt_model.named_parameters() if ('visual_backbone' in name and 'pooling_head' not in name)]
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n not in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n not in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and n in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                        "lr": self.args.vision_lr,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and n in vision_backbone_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                        "lr": self.args.vision_lr,
                    },
                ]
            else:
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                ]

            if self.optimizer_cls_and_kwargs is not None:
                optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
            else:
                optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model)

            # Overwrite `params` in case it's created by `get_optimizer_cls_and_kwargs`
            # e.g. for GaLore optimizer.
            if "params" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("params")

            # Overwrite `model` in case it's created by `get_optimizer_cls_and_kwargs`
            # e.g. for LOMO optimizer.
            if "model" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("model")

            # For layer-wise dummy optimizers we overwrite optimizer_grouped_parameters with `optimizer_dict`
            # to avoid arguments conflicts.
            if "optimizer_dict" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("optimizer_dict")

            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

            if "bitsandbytes" in str(optimizer_cls) and optimizer_kwargs.get("optim_bits", None) == 8:
                import bitsandbytes

                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()

                skipped = 0
                for module in opt_model.modules():
                    if isinstance(module, nn.Embedding):
                        skipped += sum({p.data_ptr(): p.numel() for p in module.parameters()}.values())
                        logger.info(f"skipped {module}: {skipped / 2**20}M params")
                        manager.register_module_override(module, "weight", {"optim_bits": 32})
                        logger.debug(f"bitsandbytes: will optimize {module} in fp32")
                logger.info(f"skipped: {skipped / 2**20}M params")

        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(self.optimizer)

        return self.optimizer