| | |
| | |
| | |
| | |
| |
|
| | """ |
| | 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 |
| |
|
