Use accelerator for supporting other device (bigscience-workshop#104)

* xpu support (bigscience-workshop#55)

* port accel abs interfece

* WA for run3.6b

* move on

* fix current_dievice

* fix typo

* enable to run 345M GPT

* delete apex_patch

* add TODO xpu compatible tg for xpu WA

* use deepspeed launcher

* enable run3.6b bf16

* add zero2 config json

* readd enable_each_rank_log

* fix typos

* add ccl arg

* fix

* use short word

* use no-masked-softmax-fusion

* readd

* set train  iters to 10

* remove duplicate line

* change assert msg

* update format

* add whitespace

* update path

* update note

* update

* fix typos

* delete notes

* update format

* update xpu check to cuda check

* update

* clean up file

* fix typos

* add python based gradient clipping

* change condition for python based path

Author: inkcherry

megatron/__init__.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 import os
 import torch
-
+from deepspeed.accelerator import get_accelerator
 from .package_info import (
     __description__,
     __contact_names__,
@@ -64,7 +64,7 @@ def is_rank_0():
     """Check whether it is rank 0. For AML, check if it is rank 0 of a node"""
     if torch.distributed.is_initialized():
         if torch.distributed.get_rank() == 0 or (
-            is_aml() and torch.distributed.get_rank() % torch.cuda.device_count() == 0
+            is_aml() and torch.distributed.get_rank() % get_accelerator().device_count() == 0
             ):
             return True
         else:

megatron/arguments.py

@@ -653,7 +653,7 @@ def _add_distributed_args(parser):
     group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None,
                        help='Number of layers per virtual pipeline stage')
     group.add_argument('--distributed-backend', default='nccl',
-                       choices=['nccl', 'gloo'],
+                       choices=['nccl', 'gloo', 'ccl'],
                        help='Which backend to use for distributed training.')
     group.add_argument('--DDP-impl', default='local',
                        choices=['local', 'torch'],

megatron/checkpointing.py

@@ -19,7 +19,7 @@
 import random
 import sys
 import numpy as np
-
+from deepspeed.accelerator import get_accelerator
 import torch
 
 from megatron import (get_args,
@@ -150,7 +150,7 @@ def save_checkpoint(iteration, model, optimizer, lr_scheduler):
             state_dict['random_rng_state'] = random.getstate()
             state_dict['np_rng_state'] = np.random.get_state()
             state_dict['torch_rng_state'] = torch.get_rng_state()
-            state_dict['cuda_rng_state'] = torch.cuda.get_rng_state()
+            state_dict['cuda_rng_state'] = get_accelerator().get_rng_state()
             state_dict['rng_tracker_states'] \
                 = mpu.get_cuda_rng_tracker().get_states()
 
@@ -417,7 +417,7 @@ def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load', strict=True
             random.setstate(state_dict['random_rng_state'])
             np.random.set_state(state_dict['np_rng_state'])
             torch.set_rng_state(state_dict['torch_rng_state'])
-            torch.cuda.set_rng_state(state_dict['cuda_rng_state'])
+            get_accelerator().set_rng_state(state_dict['cuda_rng_state'])
             # Check for empty states array
             if not state_dict['rng_tracker_states']:
                 raise KeyError

megatron/data/biencoder_dataset_utils.py

@@ -8,7 +8,7 @@
 from megatron.data.dataset_utils import create_masked_lm_predictions, \
                                             pad_and_convert_to_numpy
 from megatron.data.data_samplers import MegatronPretrainingSampler
-
+from deepspeed.accelerator import get_accelerator
 def make_attention_mask(source_block, target_block):
     """
     Returns a 2-dimensional (2-D) attention mask
@@ -187,7 +187,7 @@ def get_block_samples_mapping(block_dataset, title_dataset, data_prefix, num_epo
     # This should be a barrier but nccl barrier assumes
     # device_index=rank which is not the case for model
     # parallel case
-    counts = torch.cuda.LongTensor([1])
+    counts = get_accelerator().LongTensor([1])
     torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
     assert counts[0].item() == torch.distributed.get_world_size(
         group=mpu.get_data_parallel_group())

megatron/data/dataset_utils.py

@@ -33,7 +33,7 @@
 )
 from megatron.data.blendable_dataset import BlendableDataset
 from megatron.data.indexed_dataset import make_dataset as make_indexed_dataset
-
+from deepspeed.accelerator import get_accelerator
 DSET_TYPE_BERT = 'standard_bert'
 DSET_TYPE_ICT = 'ict'
 DSET_TYPE_T5  = 't5'
@@ -711,8 +711,8 @@ def get_samples_mapping(indexed_dataset,
     # This should be a barrier but nccl barrier assumes
     # device_index=rank which is not the case for model
     # parallel case
-    if torch.cuda.device_count() > 0: # Skip when CPU-only
-        counts = torch.cuda.LongTensor([1])
+    if get_accelerator().device_count() > 0: # Skip when CPU-only
+        counts = get_accelerator().LongTensor([1])
         torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
         torch.distributed.all_reduce(counts, group=mpu.get_pipeline_model_parallel_group())
         assert counts[0].item() == (

megatron/data/gpt_dataset.py

@@ -20,7 +20,7 @@
 
 import numpy as np
 import torch
-
+from deepspeed.accelerator import get_accelerator
 from megatron import mpu, is_rank_0, print_rank_0, get_args
 from megatron.data.blendable_dataset import BlendableDataset
 from megatron.data.dataset_utils import get_datasets_weights_and_num_samples
@@ -313,7 +313,7 @@ def _build_index_mappings(name, data_prefix, documents, sizes,
     # This should be a barrier but nccl barrier assumes
     # device_index=rank which is not the case for model
     # parallel case
-    counts = torch.cuda.LongTensor([1])
+    counts = get_accelerator().LongTensor([1])
     torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
     torch.distributed.all_reduce(counts, group=mpu.get_pipeline_model_parallel_group())
     assert counts[0].item() == (

megatron/data/realm_dataset_utils.py

@@ -7,7 +7,7 @@
 from megatron import mpu, print_rank_0
 from megatron.data.dataset_utils import create_masked_lm_predictions, pad_and_convert_to_numpy
 from megatron import get_args, get_tokenizer, print_rank_0, mpu
-
+from deepspeed.accelerator import get_accelerator
 
 def get_one_epoch_dataloader(dataset, micro_batch_size=None):
     """Specifically one epoch to be used in an indexing job."""
@@ -177,7 +177,7 @@ def get_block_samples_mapping(block_dataset, title_dataset, data_prefix, num_epo
     # This should be a barrier but nccl barrier assumes
     # device_index=rank which is not the case for model
     # parallel case
-    counts = torch.cuda.LongTensor([1])
+    counts = get_accelerator().LongTensor([1])
     torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
     assert counts[0].item() == torch.distributed.get_world_size(
         group=mpu.get_data_parallel_group())

megatron/global_vars.py

@@ -24,7 +24,7 @@
 from megatron.tokenizer import build_tokenizer
 from .arguments import parse_args
 from .microbatches import build_num_microbatches_calculator
-
+from deepspeed.accelerator import get_accelerator
 _GLOBAL_ARGS = None
 _GLOBAL_NUM_MICROBATCHES_CALCULATOR = None
 _GLOBAL_TOKENIZER = None
@@ -192,14 +192,14 @@ def __init__(self, name):
     def start(self):
         """Start the timer."""
         assert not self.started_, 'timer has already been started'
-        torch.cuda.synchronize()
+        get_accelerator().synchronize()
         self.start_time = time.time()
         self.started_ = True
 
     def stop(self):
         """Stop the timer."""
         assert self.started_, 'timer is not started'
-        torch.cuda.synchronize()
+        get_accelerator().synchronize()
         self.elapsed_ += (time.time() - self.start_time)
         self.started_ = False
 

megatron/initialize.py

@@ -21,7 +21,6 @@
 
 import numpy as np
 import torch
-
 from megatron import fused_kernels
 from megatron import get_adlr_autoresume
 from megatron import get_args
@@ -30,7 +29,7 @@
 from megatron.global_vars import set_global_variables
 from megatron.mpu import (set_tensor_model_parallel_rank,
                           set_tensor_model_parallel_world_size)
-
+from deepspeed.accelerator import get_accelerator
 import deepspeed
 import deepspeed.utils.groups as groups
 
@@ -46,7 +45,7 @@ def initialize_megatron(extra_args_provider=None, args_defaults={},
     """
     if not allow_no_cuda:
         # Make sure cuda is available.
-        assert torch.cuda.is_available(), 'Megatron requires CUDA.'
+        assert get_accelerator().is_available(), 'Megatron requires accelerator.'
 
     # Parse args, build tokenizer, and set adlr-autoresume,
     # tensorboard-writer, and timers.
@@ -107,7 +106,10 @@ def _compile_dependencies():
         compile_helper()
         print('>>> done with dataset index builder. Compilation time: {:.3f} '
               'seconds'.format(time.time() - start_time), flush=True)
-
+        
+    if not get_accelerator().device_name() == 'cuda':
+        print(">fused kernel is only supported in cuda, skip loading fused kernel")
+        return 
     # ==================
     # Load fused kernels
     # ==================
@@ -134,7 +136,7 @@ def _compile_dependencies():
     if _is_rank_0():
         start_time = time.time()
         print('> compiling and loading fused kernels ...', flush=True)
-        if torch.cuda.device_count() > 0: # Skip when CPU-only
+        if get_accelerator().device_count() > 0: # Skip when CPU-only
             fused_kernels.load(args)
         torch.distributed.barrier()
     else:
@@ -185,7 +187,7 @@ def setup_deepspeed_random_and_activation_checkpointing(args):
 def _initialize_distributed():
     """Initialize torch.distributed and mpu."""
     args = get_args()
-    device_count = torch.cuda.device_count()
+    device_count = get_accelerator().device_count()
     if torch.distributed.is_initialized():
 
         if args.rank == 0:
@@ -206,7 +208,7 @@ def _initialize_distributed():
             else:
                 args.local_rank = device
 
-            torch.cuda.set_device(device) # only do so when device_count > 0
+            get_accelerator().set_device(device) # only do so when device_count > 0
 
         # Call the init process
         init_method = 'tcp://'
@@ -249,14 +251,14 @@ def _set_random_seed(seed_):
     if seed_ is not None and seed_ > 0:
         # Ensure that different pipeline MP stages get different seeds.
         # No need to do so for CPU-only case.
-        if torch.cuda.device_count() == 0:
+        if get_accelerator().device_count() == 0:
             seed = seed_
         else:
             seed = seed_ + (100 * mpu.get_pipeline_model_parallel_rank())
         random.seed(seed)
         np.random.seed(seed)
         torch.manual_seed(seed)
-        if torch.cuda.device_count() > 0:
+        if get_accelerator().device_count() > 0:
             mpu.model_parallel_cuda_manual_seed(seed)
     else:
         raise ValueError('Seed ({}) should be a positive integer.'.format(seed))
@@ -284,7 +286,7 @@ def _is_rank_0():
     """Check whether it is rank 0. For AML, check if it is rank 0 of a node"""
     if torch.distributed.is_initialized():
         if torch.distributed.get_rank() == 0 or (
-            'AZUREML_EXPERIMENT_ID' in os.environ and torch.distributed.get_rank() % torch.cuda.device_count() == 0
+            'AZUREML_EXPERIMENT_ID' in os.environ and torch.distributed.get_rank() % get_accelerator().device_count() == 0
             ):
             return True
         else:

megatron/memory.py

@@ -15,7 +15,7 @@
 
 
 import torch
-
+from deepspeed.accelerator import get_accelerator
 
 # A dictionary of all the memory buffers allocated.
 _MEM_BUFFS = dict()
@@ -58,7 +58,7 @@ def __init__(self, name, numel, dtype, track_usage):
         self.dtype = dtype
         self.data = torch.empty(self.numel,
                                 dtype=self.dtype,
-                                device=torch.cuda.current_device(),
+                                device=get_accelerator().current_device_name(),
                                 requires_grad=False)
 
         # Index tracking the start of the free memory.