[BUG] OOB layer acces when using 1F1B scheduler

**Describe the bug**
When making use of the 1F1B scheduler (TP=2,PP=4,DP=1,VP=2,num_layers=24,gbatch=64,microbatch=4) results in an OOB acces on the model layer array.
Changing to VP=3 does not alleviate this issue, simply changing the index where the OOB occurs from 2 to 3.
Also, changing microbatch=16 does not change the behavior.

**To Reproduce**
I train a tiny GPT2 1.7B model and use 2 Nodes with 4 NVIDA A30s each and my training args + launch command look as follows:
```
DISTRIBUTED_ARGS=(
--nproc_per_node $SLURM_GPUS_PER_NODE
--nnodes $SLURM_JOB_NUM_NODES
--node_rank $SLURM_NODEID
--rdzv_id ${SLURM_JOB_ID:-12345}
--rdzv_backend c10d
--rdzv_endpoint $MASTER_ADDR:$MASTER_PORT
)
GPT_MODEL_ARGS+=(
--num-layers 24
--hidden-size 2304
--num-attention-heads 24
--seq-length 2048
--max-position-embeddings 2048
--attention-backend unfused
--make-vocab-size-divisible-by 1024
)
TRAINING_ARGS=(
--micro-batch-size 4
--global-batch-size 64
--train-iters 100
--weight-decay 0.1
--adam-beta1 0.9
--adam-beta2 0.95
--init-method-std 0.006
--clip-grad 1.0
--fp16
--lr 6.0e-5
--lr-decay-style cosine
--min-lr 6.0e-6
--lr-warmup-fraction .001
--lr-decay-iters 430000
--recompute-granularity full
--recompute-method uniform
--recompute-num-layers 6
--data-parallel-sharding-strategy optim
--num-layers-per-virtual-pipeline-stage 2
--microbatch-group-size-per-virtual-pipeline-stage 2
--disable-tp-comm-overlap-ag
--disable-tp-comm-overlap-rs
--disable-tp-comm-bulk-dgrad
--disable-tp-comm-bulk-wgrad
--disable-tp-comm-split-ag
--disable-tp-comm-split-rs
)
MODEL_PARALLEL_ARGS=(
--tensor-model-parallel-size 2
--pipeline-model-parallel-size 4
)
DATA_ARGS=(
--data-path $DATA_PATH
--vocab-file $VOCAB_FILE
--merge-file $MERGE_FILE
--split 949,50,1
)
EVAL_AND_LOGGING_ARGS=(
--log-memory-to-tensorboard
--log-params-norm
--log-throughput
--timing-log-level 2
--log-energy
--timing-log-option minmax
--log-interval 1
--save-interval 10000
--eval-interval 100
--save $CHECKPOINT_PATH
--eval-iters 10
--tensorboard-log-interval 1
--log-timers-to-tensorboard
--tensorboard-dir $TENSORBOARD_LOGS_PATH
)
EVAL_AND_LOGGING_ARGS+=(
--wandb-project $WANDB_PROJECT
--wandb-exp-name $WANDB_NAME
)

srun --wait=0 $CONDA_PREFIX/bin/torchrun ${DISTRIBUTED_ARGS[@]} $MEGATRON_ROOT_PATH/megatron/pretrain_gpt.py \
${GPT_MODEL_ARGS[@]} \
${TRAINING_ARGS[@]} \
${MODEL_PARALLEL_ARGS[@]} \
${DATA_ARGS[@]} \
${EVAL_AND_LOGGING_ARGS[@]}

```

**Expected behavior**
I expect the training run normally without errors if all assertions pass, which they do.

**Stack trace/logs**
- rank1 has the same traceback as rank0 (see below), all other ranks just time out
- I have added some debug statements to pring the schedule in `megatron/core/pipeline_parallel/schedules.py:forward_backward_pipelining_with_interleaving` to try and trace the issue. Their output is:
```
DEBUG: PP schedule table for 16 microbatches, 3 model chunks, and 4 microbatch group size per VP stage.
DEBUG:
PP schedule table:
[(0, 0), (1, 0), (2, 0), (3, 0), (0, 1), (1, 1), (2, 1), (3, 1), (0, 2), (1, 2), (2, 2), (3, 2), (4, 0), (5, 0), (6, 0), (7, 0), (4, 1), (5, 1), (6, 1), (7, 1), (4, 2), (5, 2), (6, 2), (7, 2), (8, 0), (9, 0), (10, 0), (11, 0), (8, 1), (9, 1), (10, 1), (11, 1), (8, 2), (9, 2), (10, 2), (11, 2), (12, 0), (13, 0), (14, 0), (15, 0), (12, 1), (13, 1), (14, 1), (15, 1), (12, 2), (13, 2), (14, 2), (15, 2)]
PP microbatch_id table:
(0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 12, 13, 14, 15, 12, 13, 14, 15, 12, 13, 14, 15)
PP model_chunk_id table:
(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2)
```
- Apart from that, there is also a traceback from rank0:
```
['Traceback (most recent call last):\n', ' File "/home/username/Project/megatron/pretrain_gpt.py", line 306, in forward_step\n output_tensor = model(\n ^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl\n return self._call_impl(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl\n return forward_call(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/distributed/data_parallel_base.py", line 22, in forward\n return self.module(*inputs, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl\n return self._call_impl(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl\n return forward_call(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/transformer/module.py", line 237, in forward\n outputs = self.module(*inputs, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl\n return self._call_impl(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl\n return forward_call(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/models/gpt/gpt_model.py", line 372, in forward\n hidden_states = self.decoder(\n ^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl\n return self._call_impl(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl\n return forward_call(*args, **kwargs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 563, in forward\n hidden_states = self._checkpointed_forward(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 433, in _checkpointed_forward\n hidden_states, context = checkpoint_handler(\n ^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 417, in checkpoint_handler\n return tensor_parallel.checkpoint(\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/tensor_parallel/random.py", line 477, in checkpoint\n return CheckpointFunction.apply(function, distribute_saved_activations, *args)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/autograd/function.py", line 575, in apply\n return super().apply(*args, **kwargs) # type: ignore[misc]\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/tensor_parallel/random.py", line 423, in forward\n outputs = run_function(*args)\n ^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 381, in custom_forward\n layer = self._get_layer(index)\n ^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 361, in _get_layer\n return self.layers[layer_number]\n ~~~~~~~~~~~^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/container.py", line 324, in __getitem__\n return self._modules[self._get_abs_string_index(idx)]\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n', ' File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/container.py", line 306, in _get_abs_string_index\n raise IndexError(f"index {idx} is out of range")\n', 'IndexError: index 2 is out of range\n']
[rank0]: Traceback (most recent call last):
[rank0]: File "/home/username/Project/megatron/pretrain_gpt.py", line 385, in
[rank0]: pretrain(
[rank0]: File "/home/username/Project/megatron/megatron/training/training.py", line 865, in pretrain
[rank0]: iteration, num_floating_point_operations_so_far = train(
[rank0]: ^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/training/training.py", line 2306, in train
[rank0]: ) = train_step(
[rank0]: ^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/training/training.py", line 1396, in train_step
[rank0]: losses_reduced = forward_backward_func(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/pipeline_parallel/schedules.py", line 1186, in forward_backward_pipelining_with_interleaving
[rank0]: output_tensor = forward_step_helper(k, microbatch_id, checkpoint_activations_microbatch)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/pipeline_parallel/schedules.py", line 1032, in forward_step_helper
[rank0]: output_tensor, num_tokens = forward_step(
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/pipeline_parallel/schedules.py", line 289, in forward_step
[rank0]: output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/pretrain_gpt.py", line 306, in forward_step
[rank0]: output_tensor = model(
[rank0]: ^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/distributed/data_parallel_base.py", line 22, in forward
[rank0]: return self.module(*inputs, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/transformer/module.py", line 237, in forward
[rank0]: outputs = self.module(*inputs, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/models/gpt/gpt_model.py", line 372, in forward
[rank0]: hidden_states = self.decoder(
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 563, in forward
[rank0]: hidden_states = self._checkpointed_forward(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 433, in _checkpointed_forward
[rank0]: hidden_states, context = checkpoint_handler(
[rank0]: ^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 417, in checkpoint_handler
[rank0]: return tensor_parallel.checkpoint(
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/tensor_parallel/random.py", line 477, in checkpoint
[rank0]: return CheckpointFunction.apply(function, distribute_saved_activations, *args)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/autograd/function.py", line 575, in apply
[rank0]: return super().apply(*args, **kwargs) # type: ignore[misc]
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/tensor_parallel/random.py", line 423, in forward
[rank0]: outputs = run_function(*args)
[rank0]: ^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 381, in custom_forward
[rank0]: layer = self._get_layer(index)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/Project/megatron/megatron/core/transformer/transformer_block.py", line 361, in _get_layer
[rank0]: return self.layers[layer_number]
[rank0]: ~~~~~~~~~~~^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/container.py", line 324, in __getitem__
[rank0]: return self._modules[self._get_abs_string_index(idx)]
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/username/.conda/envs/megatron/lib/python3.12/site-packages/torch/nn/modules/container.py", line 306, in _get_abs_string_index
[rank0]: raise IndexError(f"index {idx} is out of range")
[rank0]: IndexError: index 2 is out of range
```

**Environment (please complete the following information):**
- Megatron-LM commit: 73a938ee0cbf771ab2b3fec8e5b1d1abb7ce4333
- PyTorch version: 2.7.1+cu128 (pypi)
- CUDA version: 12.8 (spack)
- NCCL version: 2.27.6.1 (conda)

**Proposed fix**
Change the generation of the `model_chunk_id` so that the values generated stay in bounds.