GGitHubperf: reduce Megatron training memory peaks (#1555)

ea306c9c创建于 8月7日历史提交

文件	最后提交记录	最后更新时间
README.md	perf: reduce Megatron training memory peaks (#1555) * perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: #1555	1 个月前
postprocess_profile.py	perf: reduce Megatron training memory peaks (#1555) * perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: #1555	1 个月前
profile_rank_utils.py	perf: reduce Megatron training memory peaks (#1555) * perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: #1555	1 个月前
qwen3_30b_a3b_sft_profile.yaml	perf: reduce Megatron training memory peaks (#1555) * perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: #1555	1 个月前
qwen3_5_35b_a3b_sft_profile.yaml	perf: reduce Megatron training memory peaks (#1555) * perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: #1555	1 个月前
run_qwen3_30b_a3b_sft_profile.sh	perf: reduce Megatron training memory peaks (#1555) * perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: #1555	1 个月前
train_sft_profile.py	perf: reduce Megatron training memory peaks (#1555) * perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: #1555	1 个月前

Qwen3-30B-A3B SFT Profile

This directory provides an end-to-end SFT profiling example. It uses Qwen3-30B-A3B and deterministic fake 128K-token SFT data to collect kernel and memory profiles separately.

Files

train_sft_profile.py: SFT entry point. Instead of reading an external JSONL file, it uses the real tokenizer to encode a structured SWE/code-repair conversation, then repeats and truncates it to 131072 tokens.
qwen3_30b_a3b_sft_profile.yaml: Megatron MoE profiling configuration for one node with eight GPUs. By default it uses megatron:(attn:d2c4|ffn:d1e8): DP=2 and CP=4 for attention, and EP=8 for the MoE FFN.
run_qwen3_30b_a3b_sft_profile.sh: Runs kernel and memory profiles.
postprocess_profile.py: Generates kernel Chrome trace views and profile summaries.

Quick start

Run from the repository root:

MODEL_PATH=/path/to/Qwen3-30B-A3B \
FILEROOT=/path/to/shared/experiments \
bash examples/profile/run_qwen3_30b_a3b_sft_profile.sh

By default, the script profiles steps 0 and 1. With PROFILE_KINDS=kernel,memory, it runs four trials in total:

kernel: enables perf_tracer.profile_steps, starts the PyTorch profiler in the SFT training step, and produces CPU, CUDA, and GPU kernel traces.
memory: enables memory_profiler.profile_steps and produces PyTorch CUDA allocator snapshots.

Default settings:

PROFILE_STEPS=0,1
# TOTAL_STEPS is unset by default; each trial uses profile_step + 1.
# PROFILE_RANKS is unset by default; the script computes each PP stage's rank0.
PROFILE_KINDS=kernel,memory
PROFILE_FAKE_SEQ_LEN=131072
PROFILE_FAKE_DATASET_SIZE=8
PROFILE_FAKE_LOSS_START_RATIO=0.5
TRAIN_BATCH_SIZE=4
PROFILE_N_MBS=2
LOGPROBS_CHUNK_SIZE=128
LM_HEAD_LOSS_CHUNK_SIZE=0
USE_PRECISION_AWARE_OPTIMIZER=true
MAIN_GRADS_DTYPE=bfloat16
USE_DETERMINISTIC_ALGORITHMS=false
STOP_ON_FAILURE=1

Run only the kernel profile:

PROFILE_KINDS=kernel \
MODEL_PATH=/path/to/Qwen3-30B-A3B \
bash examples/profile/run_qwen3_30b_a3b_sft_profile.sh

Run only the memory profile and collect multiple ranks:

PROFILE_KINDS=memory \
PROFILE_RANKS=0,1,4-7 \
MODEL_PATH=/path/to/Qwen3-30B-A3B \
bash examples/profile/run_qwen3_30b_a3b_sft_profile.sh

When PROFILE_RANKS is unset, the script derives rank 0 of every pipeline stage from actor.backend. The default megatron:(attn:d2c4|ffn:d1e8) configuration has PP=1, so it resolves to 0. An explicit PROFILE_RANKS value overrides the computed default; set it to an empty string to collect all ranks.

Fake 128K data

The fake-data generator in train_sft_profile.py follows three rules:

It uses the target model's tokenizer to encode a structured conversation instead of filling the sequence with one token.
It repeats the same token sequence to the fixed PROFILE_FAKE_SEQ_LEN=131072, ensuring that every parallel layout sees identical token content.
By default, loss_mask is 1 from the 50% position onward. The first half is the prompt/context and the second half is the assistant target.

This avoids profile noise from real-data I/O and dynamic filtering while retaining a tokenizer distribution that includes long context, code blocks, JSON fragments, and natural-language targets.

Outputs

The script stores run-side outputs under examples/profile/ by default:

examples/profile/profile_data/<timestamp>_qwen3-30b-a3b_fake128k_sft_profile/
  profile_settings.log
  summary.tsv
  qwen3_30b_a3b_fake128k_kernel_step0_<timestamp>/
    launcher.log
    nvidia_smi.csv
    profile_summary.json
    profile_summary.md
    kernel_traces/master/
      traces-r0.chrome.json
      traces-r0.split_clean.chrome.json
      traces-r0.gpu_only.chrome.json
      traces-r0.cpu_only.chrome.json
      traces-r0.cuda_api_only.chrome.json
  qwen3_30b_a3b_fake128k_memory_step0_<timestamp>/
    launcher.log
    nvidia_smi.csv
    profile_summary.json
    profile_summary.md
    memory_snapshots/step_0/
      snapshot_*.pickle
  qwen3_30b_a3b_fake128k_kernel_step1_<timestamp>/
    launcher.log
    nvidia_smi.csv
    profile_summary.json
    profile_summary.md
    kernel_traces/master/
      traces-r0.chrome.json
      traces-r0.split_clean.chrome.json
      traces-r0.gpu_only.chrome.json
      traces-r0.cpu_only.chrome.json
      traces-r0.cuda_api_only.chrome.json
  qwen3_30b_a3b_fake128k_memory_step1_<timestamp>/
    launcher.log
    nvidia_smi.csv
    profile_summary.json
    profile_summary.md
    memory_snapshots/step_1/
      snapshot_*.pickle

The original AReaL logs remain under FILEROOT:

${FILEROOT}/logs/<user>/qwen3-30b-a3b-sft-profile/<trial_name>/
  trainer.log
  perf_tracer/<role>/traces-r*.jsonl
  memory_snapshots/step_<profile_step>/snapshot_*.pickle

Kernel-profile postprocessing generates Chrome trace views next to each trace file and copies them to profile_data/.../kernel_traces/<role>/:

traces-r0.chrome.json
traces-r0.split_clean.chrome.json
traces-r0.gpu_only.chrome.json
traces-r0.cpu_only.chrome.json
traces-r0.cuda_api_only.chrome.json

Open these .chrome.json files in Chrome at chrome://tracing or in Perfetto. With the default unset PROFILE_RANKS and PP=1 configuration, only rank 0 is collected. To compare GPU memory across DP, CP, and EP ranks, explicitly set PROFILE_RANKS= to collect all ranks.

The default 128K configuration with DP=2 and CP=4 disables deterministic algorithms. Set USE_DETERMINISTIC_ALGORITHMS=true to enable them.

Standalone postprocessing

To regenerate summaries and kernel trace views for an existing profile:

python examples/profile/postprocess_profile.py \
  --profile-kind kernel \
  --profile-step 1 \
  --log-dir /path/to/FILEROOT/logs/<user>/qwen3-30b-a3b-sft-profile/<trial_name> \
  --run-dir examples/profile/profile_data/reprocess/<trial_name> \
  --trainer-log /path/to/trainer.log \
  --nvidia-smi-csv /path/to/nvidia_smi.csv

Notes

Run kernel and memory profiles separately. PyTorch profiler memory recording disturbs the kernel timeline, so kernel trials set AREAL_TORCH_PROFILER_PROFILE_MEMORY=false by default.
This 128K example sets actor.megatron.ddp.grad_reduce_in_fp32=false to reduce the steady-state gradient-reduction memory footprint in step 1.
LOGPROBS_CHUNK_SIZE sets actor.logprobs_chunk_size, which controls the sequence chunk used for logprob and entropy computation. It defaults to 128 to reduce the steady-state temporary memory peak of the vocab-parallel path in step 1.
The Megatron precision-aware optimizer is enabled by default, and TE Adam consumes the BF16 distributed gradient buffer directly. Main parameters and both Adam moments remain FP32. This avoids allocating another FP32 gradient for every parameter at the end of the optimizer step. Set USE_PRECISION_AWARE_OPTIMIZER=false MAIN_GRADS_DTYPE=float32 to restore the original path.
PROFILE_RANKS controls which ranks are collected. When unset, the script computes rank 0 of every pipeline stage from actor.backend; an empty value selects all ranks; 0,2-4 selects ranks 0, 2, 3, and 4.
Qwen3-30B-A3B requires an eight-GPU profiling environment. Without suitable GPUs and model weights, only the scripts and postprocessing can be validated; CUDA traces and memory snapshots cannot be produced.