self.module.language_model.encoder.layers[3].mlp.deepspeed_moe.experts.deepspeed_experts[0].dense_4h_to_h.weight.flatten()[0]

Parallelism overviews

ZeRO-2 and ZeRO-3 are incompatible with out pipeline parallelism engine. ZeRO-2 partitions gradients that the pipeline engine assumes are intact. Similarly, ZeRO-3 partitions parameters that the pipeline engine assumes are intact. Note that pipeline parallelism already offers some of these advantages by partitioning the model directly, and then ZeRO-1 (with optional offload) can be combined to further partition the optimizer. (source)

• MoE doesn’t support PP because PP requires separate model code where we do not have MoE support. But it does support TP (MP). (source)

• PR-MoE

  As Phenomenon-I in Section 4.1.1 suggested that leveraging MoE at the later layers bring more benefits, our new architecture utilizes more experts in the last few layers as compared to previous layers. This gives the Pyramid-MoE design, where we show an example in Figure 3 (right)–the last two layers have 2x experts as the previous layers. Meanwhile, considering Phenomenon II, we propose the Residual-MoE architecture, where each token separately passes one fixed MLP module and one chosen expert as shown in Figure 3 (right), where orange blocks are the fixed MLP.

  

Code flows

• Megatron

  • pretrain_gpt.py
    • pretrain
      • setup_model
        • DistributedDataParallel()
          • register_hook to accumulate gradients on backward [megablocks does this only for non-experts]
      • train
        • train_step in loop
          • forward_backward_func, which can be forward_backward_no_pipelining, forward_backward_pipelining_with_interleaving, forward_backward_pipelining_without_interleaving →
  • → forward_backward_no_pipelining or with pipelining (if not DS pipelining!)
    • forward_step (schedules.py)
      • forward_step (pretrain_gpt.py)
        • GPTModel
          • TransformerLanguageModel.forward
            • ParallelTransformer.forward
              • ParallelTransformer.checkpointed_forward
                • for each layer
                  • TP.checkpoint
                    • ParallelTransformerLayer.forward
                      • mlp/moe
                  • extend moe_losses
                • return hidden_states, moe_losses
        • return output, partial(loss_func, summed moe loss)
    • loss_func(output)
  • → or with pipelining (in DS)
    • forward_step (schedules.py)
      • forward_step
        • PipelineModule.forward superclass of GPTModelPipe [note how this skips directly to the layers, unlike GPTModel above]
          • …
          • ParallelTransformerLayerPipe.forward →

• gpt-neox

  • train.py
    • pretrain
      • setup_model_and_optimizer
        • get_optimizer
          • get_params_for_weight_decay_optimization()
          • default: apex Adam or else deepspeed Adam
        • model = GPT2ModelPipe.ctor
          • PipelineModule.ctor(loss_fn=cross_entropy)
        • deepspeed.initialize
      • train
        • train_step
          • if pipeline parallel: train_step_pipe
            • PipelineEngine.train_batch → (see DS)
          • else:
            • forward_step
              • loss = cross_entropy(…)
  • eval.py
    • setup_for_inference_or_eval
      • setup_model_and_optimizer →

• DS

  • user code: model = PipelineModule.ctor

    • ctor: given list of LayerSpecs (not layers!)
      • _partition_layers
        • _set_bounds: set _local_start/stop
      • _build
        • make layers from layerspecs, from local start to stop only

  • user code: MoE(expert=SomeModelToClone())

    • …

      • Experts() deep-clones N times and marks each with .allreduce/.group_name

        class Experts(torch.nn.Module):
        
            def __init__(self, expert, num_local_experts=1, expert_group_name=None):
                super(Experts, self).__init__()
        
                self.deepspeed_experts = torch.nn.ModuleList([copy.deepcopy(expert) for i in range(num_local_experts)])
                self.num_local_experts = num_local_experts
        
                # TODO: revisit allreduce for moe.gate...
                for expert in self.deepspeed_experts:
                    # TODO: Create param groups to handle expert + data case (e.g. param.group = moe_group)
                    for name, param in expert.named_parameters():
                        param.allreduce = False
                        param.group_name = expert_group_name
        
        

  • user code: model = deepspeed.initialize(model)

    • PipelineEngine() or DeepSpeedEngine() →
      • PipelineEngine only accepts PipelineModule models
      • configure_distributed_model
        • MoE.set_deepspeed_parallelism
      • if optimizer: self._configure_optimizer(optimizer, params)
        • if zero: self._configure_zero_optimizer →
      • if no optimizer but zero: self._configure_zero_optimizer()
        • DeepSpeedZeroOptimizer()
          • self._configure_moe_settings()
            • MoE.set_deepspeed_parallelism()
              • MoE._create_process_group()
                • _create_expert_and_data_parallel, or
                • _create_expert_data_and_model_parallel
    • → DeepSpeedEngine
      • _configure_distributed_model
        • set_deepspeed_parallelism
          • _create_process_groups
            • _get_expert_parallel_group
            • sharded_moe._set_ep_group

  • user code:

    • PipelineEngine.train_batch

      # Do the work
      sched = schedule.TrainSchedule(micro_batches=self.micro_batches,
                                     stages=self.num_stages,
                                     stage_id=self.stage_id)
      self._exec_schedule(sched)
      self.agg_train_loss = self._aggregate_total_loss()
      

      • TrainSchedule.steps generates the “instructions” schedule. Maps pipeline step → micro-batch, and which direction.

        • instructions are

              _INSTRUCTION_MAP = {
                  schedule.OptimizerStep: _exec_optimizer_step,
                  schedule.ReduceGrads: _exec_reduce_grads,
                  schedule.ReduceTiedGrads: _exec_reduce_tied_grads,
                  schedule.LoadMicroBatch: _exec_load_micro_batch,
                  schedule.ForwardPass: _exec_forward_pass,
                  schedule.BackwardPass: _exec_backward_pass,
                  schedule.SendActivation: _exec_send_activations,
                  schedule.RecvActivation: _exec_recv_activations,
                  schedule.SendGrad: _exec_send_grads,
                  schedule.RecvGrad: _exec_recv_grads,
              }
          

        • _exec_forward_pass

          • DeepspeedEngine.forward
            • loss = module.forward (PipelineModule.forward)

        • has its own p2p comm system

  • DSE.backward

  • writing _EXPERT_DATA_PARALLEL_GROUP:

Misc notes

• How does DP work in neox?
• How does moe loss get calculated, in megablocks? I don’t see where it is actually getting communicated (across pipeline). Or is loss calculated for each pipeline stage? (Seems different from how it works in megatron / in neox?)
• How should expert parallelism work in neox? Should we introduce a new grouping for expert parallelism? (Did I already do that?)
• How does router stay in sync across experts, but experts vary?

Parallelism details