Only local rank zero is trained in piped parallel(DDPSequentialPlugin).
🐛 Only local rank zero is trained in piped parallel
The 1.2B model was being trained, but the performance did not improve anymore, so I checked a little more, and only the rank zero model parameter was being trained.
This is what I checked.
# initial model state_dict
d = torch.load('kogpt2_from_pretrained_210118_new_format')

# after 20 steps from d
b = torch.load('kogpt2_from_pretrained_210118_test_3')

for k, v in d.items():
    # print(k)
    if torch.all(v == b[k]).item() == False:
        print(k)
0.weight
2.ln_1.weight
2.ln_1.bias
2.ln_2.weight
2.ln_2.bias
2.attn.in_proj_weight
2.attn.in_proj_bias
2.attn.out_proj.weight
2.attn.out_proj.bias
2.mlp.c_fc.weight
2.mlp.c_fc.bias
2.mlp.c_proj.weight
2.mlp.c_proj.bias
3.ln_1.weight
3.ln_1.bias
3.ln_2.weight
3.ln_2.bias
3.attn.in_proj_weight
3.attn.in_proj_bias
3.attn.out_proj.weight
3.attn.out_proj.bias
3.mlp.c_fc.weight
3.mlp.c_fc.bias
3.mlp.c_proj.weight
3.mlp.c_proj.bias

# after 600 steps from d
f = torch.load('kogpt2_from_pretrained_210118_test_4')

for k, v in d.items():
    # print(k)
    if torch.all(v == f[k]).item() == False:
        print(k)
0.weight
2.ln_1.weight
2.ln_1.bias
2.ln_2.weight
2.ln_2.bias
2.attn.in_proj_weight
2.attn.in_proj_bias
2.attn.out_proj.weight
2.attn.out_proj.bias
2.mlp.c_fc.weight
2.mlp.c_fc.bias
2.mlp.c_proj.weight
2.mlp.c_proj.bias
3.ln_1.weight
3.ln_1.bias
3.ln_2.weight
3.ln_2.bias
3.attn.in_proj_weight
3.attn.in_proj_bias
3.attn.out_proj.weight
3.attn.out_proj.bias
3.mlp.c_fc.weight
3.mlp.c_fc.bias
3.mlp.c_proj.weight
3.mlp.c_proj.bias
plugin = DDPSequentialPlugin(balance=[4, 9, 9, 5], microbatches=64,
                             rpc_timeout_sec=args.rpc_timeout_sec)
python train.py .... --accelerator ddp --use_ddp_sequential --rpc_timeout_sec 600 --gpus 4 --gradient_clip_val 1.0 --accumulate_grad_batches 3 
To Reproduce
SeqGPT2Model() returns nn.Sequencial.
class GPT2(LightningModule):
    def __init__(self, hparams, **kwargs) -> None:
        super(GPT2, self).__init__(hparams, **kwargs)
        self.sequential_module = SeqGPT2Model(GPT2Cfg1B)
        self.loss_function = torch.nn.CrossEntropyLoss(
            ignore_index=model_arch['pad_token_id'])

    def forward(self, batch):
        logits = self.sequential_module(batch['decoder_in'])
        return  logits

    def training_step(self, batch, batch_idx):
        label = batch['label']
        y_hat = self.forward(batch)
        train_loss = self.loss_function(y_hat.transpose(1, 2), label)
        self.log('train_loss', train_loss)
        return train_loss

Are there any points I am missing from model training?
Expected behavior
Model parameters of all ranks must be trained.
Environment

PyTorch Version (e.g., 1.0): 1.6.0
OS (e.g., Linux): CentOS 7
How you installed PyTorch (conda, pip, source): conda
Build command you used (if compiling from source):
Python version: 3.7
CUDA/cuDNN version: 10.2
GPU models and configuration: V100
Any other relevant information:

Additional context