Synthetic tabular data generation non-reproducibility
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

Problem with fitting model step-by-step in [this](https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Megatron_Synthetic_Tabular_Data_Generation.ipynb) notebook. 

**Steps/Code to reproduce bug**

On step
```
!python megatron_gpt_pretraining.py \
        trainer.devices={NUM_GPUS} \
        trainer.accelerator=gpu \
        trainer.log_every_n_steps=100 \
        trainer.val_check_interval=500 \
        trainer.accumulate_grad_batches=1 \
        trainer.max_steps=10000 \
        trainer.precision=16 \
        trainer.gradient_clip_val=1.0 \
        exp_manager.exp_dir=gpt_creditcard_results \
        model.tensor_model_parallel_size={TENSOR_MP_SIZE} \
        model.pipeline_model_parallel_size={PIPELINE_MP_SIZE} \
        model.optim.name=fused_adam \
        model.optim.lr=2e-4 \
        model.optim.sched.warmup_steps=2 \
        model.optim.sched.constant_steps=2 \
        model.optim.sched.min_lr=8e-5 \
        model.max_position_embeddings={SEQ_LENGTH} \
        model.encoder_seq_length={SEQ_LENGTH} \
        model.data.seq_length={SEQ_LENGTH} \
        model.tokenizer.type=Tabular \
        model.tokenizer.library=tabular \
        model.tokenizer.vocab_file={CC_OUTPUT_P} \
        model.tokenizer.delimiter=\',\' \
        model.data.eod_mask_loss=True \
        model.data.splits_string=\'3800,198,2\' \
        model.num_layers={NUM_LAYERS} \
        model.hidden_size={HIDDEN_SIZE} \
        model.num_attention_heads={NUM_ATTENTION_HEADS} \
        model.activations_checkpoint_method=\'block\' \
        model.activations_checkpoint_num_layers=1 \
        model.data.data_prefix=[tabular_data_text_document]

```

I've got such traceback
```
Traceback (most recent call last):
  File "megatron_gpt_pretraining.py", line 88, in main
    trainer.fit(model)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 770, in fit
    self._call_and_handle_interrupt(
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 721, in _call_and_handle_interrupt
    return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 811, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1236, in _run
    results = self._run_stage()
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1323, in _run_stage
    return self._run_train()
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1345, in _run_train
    self._run_sanity_check()
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1413, in _run_sanity_check
    val_loop.run()
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 154, in advance
    dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 128, in advance
    output = self._evaluation_step(**kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 226, in _evaluation_step
    output = self.trainer._call_strategy_hook("validation_step", *kwargs.values())
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1765, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 355, in validation_step
    return self.model(*args, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 888, in forward
    output = self.module(*inputs, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 93, in forward
    return self.module.validation_step(*inputs, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 397, in validation_step
    losses_reduced_per_micro_batch = forward_backward_no_pipelining(
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py", line 90, in forward_backward_no_pipelining
    output_tensor = forward_step(
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/apex/transformer/pipeline_parallel/schedules/common.py", line 299, in forward_step
    output_tensor, loss_func = forward_step_func(batch, model)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 330, in fwd_output_and_loss_func
    output_tensor = model(tokens, position_ids, attention_mask, labels)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron/gpt_model.py", line 173, in forward
    lm_output = self.language_model(
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/language_model.py", line 505, in forward
    encoder_output = self.encoder(
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/transformer.py", line 1757, in forward
    hidden_states = self._checkpointed_forward(
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/transformer.py", line 1675, in _checkpointed_forward
    hidden_states = tensor_parallel.checkpoint(
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/apex/transformer/tensor_parallel/random.py", line 311, in checkpoint
    return CheckpointFunction.apply(function, distribute_saved_activations, *args)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/apex/transformer/tensor_parallel/random.py", line 256, in forward
    outputs = run_function(*args)
  File "/home/user/conda/envs/nemo/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/transformer.py", line 1626, in custom_forward
    encoder_decoder_position_bias = inputs[6]
IndexError: tuple index out of range
```

All previous cells worked fine.

**Expected behavior**

Model fitting process

**Environment overview (please complete the following information)**

 - Environment location: [Cloud provider]
 - Method of NeMo install: [pip install nemo_toolkit['all']]. 


**Environment details**

- OS version Ubuntu 18.04.6 LTS
- PyTorch 1.10.2+cu111
- Python 3.8