# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
Train a network across multiple GPUs.
"""

from fairseq.dataclass.configs import FairseqConfig
from fairseq.distributed import utils as distributed_utils
from fairseq.trainer import Trainer

try:
    from fairseq.model_parallel.megatron.mpu import (
        get_data_parallel_rank,
        get_data_parallel_world_size,
        get_model_parallel_src_rank,
        get_cuda_rng_tracker,
    )

    has_megatron_submodule = True
except (ImportError, ModuleNotFoundError):
    has_megatron_submodule = False


class MegatronTrainer(Trainer):
    """Main class for model parallel with data parallel training."""

    def __init__(self, cfg: FairseqConfig, task, model, criterion, **kwargs):
        if not has_megatron_submodule:
            raise ImportError(
                "\n\nPlease install the megatron submodule:"
                "\n\n  git submodule update --init "
                "fairseq/model_parallel/megatron"
            )
        super().__init__(cfg, task, model, criterion, **kwargs)

    def clip_grad_norm(self, clip_norm):
        def _aggregate_model_parallel_grad_norm(total_norm):
            total_norm = total_norm ** 2
            distributed_utils.all_reduce(
                total_norm, group=distributed_utils.get_model_parallel_group()
            )
            total_norm = total_norm ** 0.5
            return total_norm

        return self.optimizer.clip_grad_norm(
            clip_norm,
            aggregate_norm_fn=_aggregate_model_parallel_grad_norm,
        )

    def save_checkpoint(self, filename, extra_state):
        """Save all training state in a checkpoint file."""
        extra_state['rng_tracker_states'] \
            = get_cuda_rng_tracker().get_states()
        super().save_checkpoint(filename, extra_state)

    def load_checkpoint(
        self,
        filename,
        reset_optimizer=False,
        reset_lr_scheduler=False,
        optimizer_overrides=None,
        reset_meters=False,
    ):
        extra_state = super().load_checkpoint(filename, reset_optimizer=reset_optimizer, reset_lr_scheduler=reset_lr_scheduler, optimizer_overrides=optimizer_overrides, reset_meters=reset_meters)
        if extra_state is not None and 'rng_tracker_states' in extra_state:
            get_cuda_rng_tracker().set_states(
                extra_state['rng_tracker_states'])
        return extra_state