wehub-resource-sync
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
# 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
|