Using Full recompute does not decrease vram usage

[JavaZeroo]

I don't understand why vram hasn't changed anything. But from the speed of training it looks as if full recompute has been enabled

baseline

TP=8
PP=1
DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"
GPT_ARGS="
    --use-mcore-models \
    --tensor-model-parallel-size ${TP} \
    --pipeline-model-parallel-size ${PP} \
    --num-layers 24 \
    --hidden-size 6144 \
    --ffn-hidden-size 16384 \
    --num-attention-heads 48 \
    --group-query-attention \
    --num-query-groups 8 \
    --seq-length 128 \
    --max-position-embeddings 128 \
    --micro-batch-size 1 \
    --global-batch-size 1 \
    --make-vocab-size-divisible-by 1 \
    --lr 1e-6 \
    --lr-decay-style constant \
    --rotary-base 1000000 \
    --train-iters 101 \
    --untie-embeddings-and-output-weights \
    --disable-bias-linear \
    --init-method-std 0.01 \
    --attention-dropout 0.0 \
    --hidden-dropout 0.0 \
    --clip-grad 1.0 \
    --adam-beta1 0.9 \
    --adam-beta2 0.95 \
    --initial-loss-scale 1 \
    --no-gradient-accumulation-fusion \
    --use-flash-attn \
    --use-rotary-position-embeddings \
    --position-embedding-type rope \
    --normalization RMSNorm \
    --norm-epsilon 1e-5 \
    --swiglu \
    --no-masked-softmax-fusion \
    --attention-softmax-in-fp32 \
    --no-load-optim \
    --no-load-rng \
    --seed 42 \
    --bf16 \
    --dist-ckpt-strictness log_unexpected \
    --no-masked-softmax-fusion \
    --no-bias-swiglu-fusion \
    --no-bias-dropout-fusion \
    --no-rope-fusion \
    --attention-softmax-in-fp32 \
    --accumulate-allreduce-grads-in-fp32 \
"

vram usage

    PID USER DEV     TYPE  GPU        GPU MEM    CPU  HOST MEM Command                                                                                                                           
 121397 root   0  Compute N/A   25932MiB  63%   102%   2036MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 121398 root   1  Compute N/A   25932MiB  63%   102%   2055MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 121399 root   2  Compute N/A   25932MiB  63%   100%   2036MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 121400 root   3  Compute N/A   25932MiB  63%   100%   2036MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 121401 root   4  Compute N/A   25932MiB  63%   102%   2037MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 121402 root   5  Compute N/A   25932MiB  63%   102%   2036MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 121403 root   6  Compute N/A   25932MiB  63%   102%   2037MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 121404 root   7  Compute N/A   25932MiB  63%   100%   2039MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers

training log

2025-03-08 03:39:44] iteration       93/     101 | consumed samples:           93 | elapsed time per iteration (ms): 283.3 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.598389E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:45] iteration       94/     101 | consumed samples:           94 | elapsed time per iteration (ms): 285.1 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.579762E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:45] iteration       95/     101 | consumed samples:           95 | elapsed time per iteration (ms): 287.5 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.523883E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:45] iteration       96/     101 | consumed samples:           96 | elapsed time per iteration (ms): 282.4 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.402811E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:46] iteration       97/     101 | consumed samples:           97 | elapsed time per iteration (ms): 286.0 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.356245E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:46] iteration       98/     101 | consumed samples:           98 | elapsed time per iteration (ms): 287.5 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.263113E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:46] iteration       99/     101 | consumed samples:           99 | elapsed time per iteration (ms): 290.5 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.263113E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:46] iteration      100/     101 | consumed samples:          100 | elapsed time per iteration (ms): 286.2 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.197920E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:39:47] iteration      101/     101 | consumed samples:          101 | elapsed time per iteration (ms): 285.7 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 2.151354E-07 | loss scale: 1.0 | grad norm: 0.000 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |

recompute

after adding:

    --recompute-granularity full \
    --recompute-method uniform \
    --recompute-num-layers 24 \

training log

 [2025-03-08 03:37:35] iteration       93/     101 | consumed samples:           93 | elapsed time per iteration (ms): 372.5 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.931662E+00 | loss scale: 1.0 | grad norm: 106.125 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:36] iteration       94/     101 | consumed samples:           94 | elapsed time per iteration (ms): 378.3 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.905890E+00 | loss scale: 1.0 | grad norm: 105.842 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:36] iteration       95/     101 | consumed samples:           95 | elapsed time per iteration (ms): 373.4 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.883155E+00 | loss scale: 1.0 | grad norm: 105.635 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:36] iteration       96/     101 | consumed samples:           96 | elapsed time per iteration (ms): 379.1 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.860850E+00 | loss scale: 1.0 | grad norm: 105.362 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:37] iteration       97/     101 | consumed samples:           97 | elapsed time per iteration (ms): 383.2 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.842552E+00 | loss scale: 1.0 | grad norm: 105.141 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:37] iteration       98/     101 | consumed samples:           98 | elapsed time per iteration (ms): 376.8 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.819790E+00 | loss scale: 1.0 | grad norm: 104.850 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:37] iteration       99/     101 | consumed samples:           99 | elapsed time per iteration (ms): 377.7 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.797119E+00 | loss scale: 1.0 | grad norm: 104.588 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:38] iteration      100/     101 | consumed samples:          100 | elapsed time per iteration (ms): 375.4 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.776598E+00 | loss scale: 1.0 | grad norm: 104.269 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |
 [2025-03-08 03:37:38] iteration      101/     101 | consumed samples:          101 | elapsed time per iteration (ms): 367.8 | learning rate: 1.000000E-06 | global batch size:     1 | lm loss: 1.756780E+00 | loss scale: 1.0 | grad norm: 104.023 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |

vram usage

   PID USER DEV     TYPE  GPU        GPU MEM    CPU  HOST MEM Command                                                                                                                           
 116035 root   7  Compute N/A   24438MiB  60%    32%   3248MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 116028 root   0  Compute N/A   24722MiB  60%    45%   3249MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 116032 root   4  Compute N/A   24726MiB  60%    26%   3251MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 116034 root   6  Compute N/A   24726MiB  60%    42%   3248MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 116030 root   2  Compute N/A   24918MiB  61%    39%   3251MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 116031 root   3  Compute N/A   25206MiB  62%    27%   3272MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 116029 root   1  Compute N/A   25210MiB  62%    31%   3270MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers
 116033 root   5  Compute N/A   25210MiB  62%    30%   3268MiB /usr/bin/python -u pretrain_gpt.py --use-mcore-models --tensor-model-parallel-size 8 --pipeline-model-parallel-size 1 --num-layers