[QUESTION] AssertionError: decoupled_learning_rate is None during logging (Llama 3 8B, PP=4, Megatron-core v0.14.0)

I am attempting to pre-train the Llama 3 8B model using the provided script under `examples/llama/train_llama3_8b_h100_fp8.sh` in Megatron-LM-core v0.14.0. The training fails immediately with an AssertionError related to the decoupled_learning_rate variable.
The error traceback shows the assert failing in the logging function (`megatron/training/training.py`, line 1555), suggesting that decoupled_learning_rate is None on a Pipeline Parallelism Rank where it is expected to be initialized.

I am using a multi-node setup with 4 servers and a total of 16 GPUs (4 GPUs per node). The full training script used is provided below:
```
#!/bin/bash

# Environment variables for performance tuning
export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
#export LOG_LEVEL=${LOG_LEVEL:-INFO}
#export NCCL_IB_TIMEOUT=${NCCL_IB_TIMEOUT:-19}
#export NVTE_FWD_LAYERNORM_SM_MARGIN=${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}
#export NVTE_BWD_LAYERNORM_SM_MARGIN=${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}
#export NCCL_P2P_NET_CHUNKSIZE=${NCCL_P2P_NET_CHUNKSIZE:-2097152}
#export NCCL_AVOID_RECORD_STREAMS=${NCCL_AVOID_RECORD_STREAMS:-1}

CHECKPOINT_PATH=${1:-"checkpoints/llama3_8b_fp8"}
TENSORBOARD_LOGS_PATH=${2:-"tensorboard_logs/llama3_8b_fp8"}
TOKENIZER_ARG=${3:-"MOCK"} # Path to tokenizer model, or "MOCK"
DATA_ARG=${4:-"MOCK"} # Data prefix, or "MOCK"

# Create directories if they don't exist
mkdir -p "$(dirname "$CHECKPOINT_PATH")"
mkdir -p "$(dirname "$TENSORBOARD_LOGS_PATH")"

# Distributed training setup
GPUS_PER_NODE=4
NUM_NODES=4
MASTER_ADDR=10.229.130.145
MASTER_PORT=29501
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))

# Path to the pretrain_gpt.py script, assuming this script is run from the root of the Megatron-LM repository
PRETRAIN_SCRIPT_PATH="pretrain_gpt.py"

# Fixed model and training parameters
TP_SIZE=2
CP_SIZE=1
PP_SIZE=4
MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=32
NUM_LAYERS=32
DTYPE="fp16"
SEQ_LENGTH=8192
MAX_POSITION_EMBEDDINGS=8192

# Data cache path (useful for both mock and real data)
DATA_CACHE_PATH="${PWD}/benchmark_cache_llama3_8b_fp8"
mkdir -p "$DATA_CACHE_PATH"

DISTRIBUTED_ARGS=(
--nproc_per_node $GPUS_PER_NODE
--nnodes $NUM_NODES
--node_rank $NODE_RANK
--master_addr $MASTER_ADDR
--master_port $MASTER_PORT
)

MODEL_ARGS=(
--use-mcore-models
--num-layers $NUM_LAYERS
--hidden-size 4096
--ffn-hidden-size 14336
--num-attention-heads 32
--group-query-attention
--num-query-groups 8
--kv-channels 128
--seq-length $SEQ_LENGTH
--max-position-embeddings $MAX_POSITION_EMBEDDINGS
--position-embedding-type rope
--rotary-base 1000000
--rotary-percent 1.0
--attention-dropout 0.0
--hidden-dropout 0.0
--swiglu
--init-method-std 0.0134
--attention-backend fused
--apply-layernorm-1p
--untie-embeddings-and-output-weights
--disable-bias-linear
)

TRAINING_ARGS=(
--micro-batch-size $MICRO_BATCH_SIZE
--global-batch-size $GLOBAL_BATCH_SIZE
# --train-samples 1953125000
# --lr-decay-samples 1949218748
# --lr-warmup-samples 3906252
--train-iters 10
--lr-decay-iters 3
--lr-warmup-iters 2
--lr 0.00015
--min-lr 0.00001
--decoupled-lr 5.0e-4 # Specific to decoupled AdamW, ensure optimizer is compatible
--decoupled-min-lr 4.5e-5 # Specific to decoupled AdamW
--lr-decay-style cosine
--clip-grad 1.0
--weight-decay 0.1
--adam-beta1 0.9
--adam-beta2 0.95
--bf16
--grad-reduce-in-bf16
--cross-entropy-loss-fusion
--calculate-per-token-loss
--manual-gc
--empty-unused-memory-level 1
--exit-duration-in-mins 235
)

# Conditional arguments based on DTYPE (FP8)
DTYPE_ARGS=()
if [[ "$DTYPE" == "fp8" ]]; then
DTYPE_ARGS+=(
"--fp8-format hybrid"
"--fp8-amax-history-len 1024"
"--fp8-amax-compute-algo max"
"--fp8-param-gather"
)
fi

# Model parallelism arguments
MODEL_PARALLEL_ARGS=(
--tensor-model-parallel-size $TP_SIZE
--context-parallel-size $CP_SIZE
--pipeline-model-parallel-size $PP_SIZE # Not explicitly set in llama script options, assume 1 if not multi-node PP
--sequence-parallel # Always enable sequence parallelism with TP_SIZE=2
)

# Distributed Data Parallel (DDP) arguments
# From original script's ddp_args
DDP_ARGS=(
--use-distributed-optimizer
--overlap-grad-reduce
--overlap-param-gather
)
TRAINING_ARGS+=("${DDP_ARGS[@]}")

# Data arguments (conditional for mock vs real data)
DATA_ARGS_LIST=()
if [[ "$TOKENIZER_ARG" == "MOCK" ]] || [[ "$DATA_ARG" == "MOCK" ]] || [[ -z "$TOKENIZER_ARG" ]]; then
DATA_ARGS_LIST+=(
"--mock-data"
"--tokenizer-type NullTokenizer"
"--vocab-size 128256"
"--data-cache-path ${DATA_CACHE_PATH}"
"--tiktoken-pattern v2"
"--split '99,1,0'"
"--no-create-attention-mask-in-dataloader"
"--no-mmap-bin-files"
"--num-workers 1"
)
else
# Settings for real data
DATA_ARGS_LIST+=(
"--data-path $DATA_ARG"
"--tokenizer-type HuggingFaceTokenizer"
"--tokenizer-model $TOKENIZER_ARG"
"--data-cache-path ${DATA_CACHE_PATH}"
"--split '99,1,0'"
"--no-create-attention-mask-in-dataloader"
"--no-mmap-bin-files"
"--num-workers 1"
# Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit.
"--vocab-size 128256"
)
fi

EVAL_AND_LOGGING_ARGS=(
--log-interval 1
--eval-iters 4
--eval-interval 200
--save-interval 1000
--log-throughput
--profile
--profile-step-start 5
--profile-step-end 7
--ckpt-format torch_dist
--distributed-timeout-minutes 60
# --save "$CHECKPOINT_PATH"
# --load "$CHECKPOINT_PATH"
# --tensorboard-dir "$TENSORBOARD_LOGS_PATH"
)

# Ensure pretrain_gpt.py is found
if [ ! -f "$PRETRAIN_SCRIPT_PATH" ]; then
echo "Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH"
echo "Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present."
exit 1
fi

# Run the training command
torchrun ${DISTRIBUTED_ARGS[@]} \
"$PRETRAIN_SCRIPT_PATH" \
${MODEL_ARGS[@]} \
${TRAINING_ARGS[@]} \
${DTYPE_ARGS[@]} \
${MODEL_PARALLEL_ARGS[@]} \
${DATA_ARGS_LIST[@]} \
${EVAL_AND_LOGGING_ARGS[@]}

set +x
```

The training command failed as follows:
```
[rank15]: Traceback (most recent call last):
[rank15]: File "/workspace/megatron/pretrain_gpt.py", line 396, in
[rank15]: pretrain(
[rank15]: File "/workspace/megatron/megatron/training/training.py", line 710, in pretrain
[rank15]: iteration, num_floating_point_operations_so_far = train(
[rank15]: ^^^^^^
[rank15]: File "/workspace/megatron/megatron/training/training.py", line 2195, in train
[rank15]: report_memory_flag = training_log(
[rank15]: ^^^^^^^^^^^^^
[rank15]: File "/workspace/megatron/megatron/training/training.py", line 1555, in training_log
[rank15]: assert decoupled_learning_rate is not None
[rank15]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank15]: AssertionError
[rank14]: Traceback (most recent call last):
[rank14]: File "/workspace/megatron/pretrain_gpt.py", line 396, in
[rank14]: pretrain(
[rank14]: File "/workspace/megatron/megatron/training/training.py", line 710, in pretrain
[rank14]: iteration, num_floating_point_operations_so_far = train(
[rank14]: ^^^^^^
[rank14]: File "/workspace/megatron/megatron/training/training.py", line 2195, in train
[rank14]: report_memory_flag = training_log(
[rank14]: ^^^^^^^^^^^^^
[rank14]: File "/workspace/megatron/megatron/training/training.py", line 1555, in training_log
[rank14]: assert decoupled_learning_rate is not None
[rank14]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank14]: AssertionError
W1106 06:30:02.458000 3410 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 3475 closing signal SIGTERM
W1106 06:30:02.460000 3410 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 3476 closing signal SIGTERM
W1106 06:30:02.461000 3410 torch/distributed/elastic/multiprocessing/api.py:900] Sending process 3477 closing signal SIGTERM
E1106 06:30:06.435000 3410 torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 3 (pid: 3478) of binary: /usr/bin/python
Traceback (most recent call last):
File "/usr/local/bin/torchrun", line 33, in
sys.exit(load_entry_point('torch==2.7.0a0+79aa17489c.nv25.4', 'console_scripts', 'torchrun')())
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 892, in main
run(args)
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 883, in run
elastic_launch(
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 139, in __call__
return launch_agent(self._config, self._entrypoint, list(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 270, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
pretrain_gpt.py FAILED
------------------------------------------------------------
Failures:

------------------------------------------------------------
Root Cause (first observed failure):
[0]:
time : 2025-11-06_06:30:02
host : Host65
rank : 15 (local_rank: 3)
exitcode : 1 (pid: 3478)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
```
**Question**: Given that the error occurs when checking for the decoupled_learning_rate, and I'm using the Decoupled AdamW specific arguments, how can I explicitly confirm or set the Decoupled AdamW optimizer in the Megatron-LM-core v0.14.0 framework?
```
--decoupled-lr 5.0e-4 # Specific to decoupled AdamW, ensure optimizer is compatible
--decoupled-min-lr 4.5e-5 # Specific to decoupled AdamW
```