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)

Short Name	Flexible Parallelism Configurations	Benefit
E	Expert	Scales the model size by increasing the number of experts
E + D	Expert + Data	Accelerates training throughput by scaling to multiple data parallel groups
E + Z	Expert + ZeRO-powered data	Partitions the nonexpert parameters to support larger base models
E + D + M	Expert + Data + Model	Supports massive hidden sizes and even larger base models than E+Z
E + D + Z	Expert + Data + ZeRO-powered data	Supports massive hidden sizes and even larger base models than E+Z
E + Z-Off + M	Expert + ZeRO-Offload + Model	Leverages both GPU and CPU memory for large MoE models on limited # of GPUs

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 →

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)
    - forward_funcs
has its own p2p comm system