megatron-memory-estimator

Name: megatron-memory-estimator
Author: yzlnew

13installs10forks

How to install megatron-memory-estimator?

yzlnew/infra-skills/megatron-memory-estimator

$npx -y skills add yzlnew/infra-skills --skill megatron-memory-estimator

Installs into the current project.

Summary

Estimate GPU memory usage for Megatron-based MoE (Mixture of Experts) and dense models. Use when users need to (1) estimate memory from HuggingFace model configs (DeepSeek-V3, Qwen, etc.), (2) plan GPU resource allocation for training, (3) compare different parallelism strategies

Files · 1

View on GitHub

SKILL.md

1	# Megatron Memory Estimator
2
3	Estimate GPU memory usage for Megatron-based models directly from HuggingFace configs or custom specifications.
4
5	## Quick Start
6
7	### Option 1: From HuggingFace Model (Recommended)
8
9	Estimate directly from HuggingFace model paths:
10
11	```bash
12	# DeepSeek-V3 (61 layers, requires layer distribution when pp>1)
13	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
14	--tp 4 --pp 4 --ep 8 --num-gpus 128 --num-layers-in-last-pipeline-stage 16
15
16	# Qwen 3
17	python scripts/estimate_from_hf.py Qwen/Qwen3-235B-A22B \
18	--tp 8 --pp 4 --ep 4 --num-gpus 128
19	```
20
21	### Option 2: From Local HF Config
22
23	```bash
24	python scripts/estimate_from_hf.py /path/to/config.json \
25	--tp 2 --pp 2 --num-gpus 8
26	```
27
28	### Option 3: Quick Parameter Testing
29
30	```bash
31	# Test different parallelism strategies
32	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
33	--tp 8 --pp 2 --ep 16 --num-layers-in-last-pipeline-stage 31 # Strategy 1 (30+31=61)
34
35	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
36	--tp 4 --pp 4 --ep 8 --num-layers-in-last-pipeline-stage 16 # Strategy 2 (15+15+15+16=61)
37
38	# Test different batch sizes
39	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
40	--tp 4 --pp 4 --ep 8 --micro-batch-size 2 --num-layers-in-last-pipeline-stage 16
41	```
42
43	## Available Scripts
44
45	### estimate_from_hf.py (Primary Script)
46
47	Automatically converts HuggingFace configs to Megatron format and estimates memory.
48
49	Key Arguments:
50	- `model_path`: HF model path or local config.json path
51	- `--tp N`: Tensor parallel size (default: 1)
52	- `--pp N`: Pipeline parallel size (default: 1)
53	- `--ep N`: Expert parallel size (default: 1, for MoE)
54	- `--cp N`: Context parallel size (default: 1)
55	- `--etp N`: Expert tensor parallel size (optional)
56	- `--vpp N`: Virtual pipeline parallel size (optional)
57	- `--micro-batch-size N`: Micro batch size (default: 1)
58	- `--seq-length N`: Sequence length (default: 4096)
59	- `--num-gpus N`: Total GPU count (default: 8)
60	- `--recompute-granularity {full,selective}`: Enable activation checkpointing
61	- `--num-layers-in-first-pipeline-stage N`: Number of layers in the first pipeline stage (use when model layers cannot be evenly divided by `--pp`)
62	- `--num-layers-in-last-pipeline-stage N`: Number of layers in the last pipeline stage (use when model layers cannot be evenly divided by `--pp`)
63	- `--verbose`: Show detailed model breakdown
64	- `--json`: Output as JSON
65
66	Examples:
67
68	```bash
69	# Basic estimation
70	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 --num-gpus 64
71
72	# With memory optimization
73	python scripts/estimate_from_hf.py Qwen/Qwen3-235B-A22B \
74	--tp 8 --pp 4 --ep 4 \
75	--recompute-granularity full \
76	--recompute-method uniform \
77	--num-gpus 128
78
79	# Verbose output
80	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
81	--tp 4 --pp 4 --ep 8 --verbose --num-layers-in-last-pipeline-stage 16
82
83	# JSON output for automation
84	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
85	--tp 4 --pp 4 --ep 8 --json --num-layers-in-last-pipeline-stage 16 > result.json
86	```
87
88
89
90	## Common Workflows
91
92	### Find Optimal Parallelism for a Model
93
94	```bash
95	# Start with model path
96	MODEL="deepseek-ai/DeepSeek-V3"
97	GPUS=128
98
99	# Test different strategies
100	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 4 --ep 8 --num-gpus $GPUS --num-layers-in-last-pipeline-stage 16
101	python scripts/estimate_from_hf.py $MODEL --tp 8 --pp 2 --ep 8 --num-gpus $GPUS --num-layers-in-last-pipeline-stage 31
102
103
104	# Choose strategy that fits GPU memory with best efficiency
105	```
106
107	### Optimize for Memory Efficiency
108
109	Progressive memory reduction:
110
111	```bash
112	# 1. Baseline
113	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 2 --num-gpus 16
114
115	# 2. Add recomputation
116	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 2 --num-gpus 16 \
117	--recompute-granularity full
118
119	# 3. Increase expert parallelism (MoE only)
120	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 2 --ep 4 --num-gpus 16 \
121	--recompute-granularity full
122
123	# 4. Increase pipeline parallelism
124	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 4 --ep 4 --num-gpus 16 \
125	--recompute-granularity full
126
127	# 5. Last resort: reduce batch size
128	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 4 --ep 4 --num-gpus 16 \
129	--recompute-granularity full --micro-batch-size 1
130	```
131
132	### Check if Model Fits Available GPUs
133
134	```bash
135	# Check if DeepSeek-V3 fits in 128x A100 80GB
136	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
137	--tp 4 --pp 4 --ep 8 --num-gpus 128 --num-layers-in-last-

People who liked this also viewed

6 picks

Type

megatron-memory-estimator

13installs10forks

View on GitHub

How to install megatron-memory-estimator?

yzlnew/infra-skills/megatron-memory-estimator

$npx -y skills add yzlnew/infra-skills --skill megatron-memory-estimator

Installs into the current project.

Summary

Files · 1

View on GitHub

SKILL.md

1	# Megatron Memory Estimator
2
3	Estimate GPU memory usage for Megatron-based models directly from HuggingFace configs or custom specifications.
4
5	## Quick Start
6
7	### Option 1: From HuggingFace Model (Recommended)
8
9	Estimate directly from HuggingFace model paths:
10
11	```bash
12	# DeepSeek-V3 (61 layers, requires layer distribution when pp>1)
13	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
14	--tp 4 --pp 4 --ep 8 --num-gpus 128 --num-layers-in-last-pipeline-stage 16
15
16	# Qwen 3
17	python scripts/estimate_from_hf.py Qwen/Qwen3-235B-A22B \
18	--tp 8 --pp 4 --ep 4 --num-gpus 128
19	```
20
21	### Option 2: From Local HF Config
22
23	```bash
24	python scripts/estimate_from_hf.py /path/to/config.json \
25	--tp 2 --pp 2 --num-gpus 8
26	```
27
28	### Option 3: Quick Parameter Testing
29
30	```bash
31	# Test different parallelism strategies
32	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
33	--tp 8 --pp 2 --ep 16 --num-layers-in-last-pipeline-stage 31 # Strategy 1 (30+31=61)
34
35	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
36	--tp 4 --pp 4 --ep 8 --num-layers-in-last-pipeline-stage 16 # Strategy 2 (15+15+15+16=61)
37
38	# Test different batch sizes
39	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
40	--tp 4 --pp 4 --ep 8 --micro-batch-size 2 --num-layers-in-last-pipeline-stage 16
41	```
42
43	## Available Scripts
44
45	### estimate_from_hf.py (Primary Script)
46
47	Automatically converts HuggingFace configs to Megatron format and estimates memory.
48
49	Key Arguments:
50	- `model_path`: HF model path or local config.json path
51	- `--tp N`: Tensor parallel size (default: 1)
52	- `--pp N`: Pipeline parallel size (default: 1)
53	- `--ep N`: Expert parallel size (default: 1, for MoE)
54	- `--cp N`: Context parallel size (default: 1)
55	- `--etp N`: Expert tensor parallel size (optional)
56	- `--vpp N`: Virtual pipeline parallel size (optional)
57	- `--micro-batch-size N`: Micro batch size (default: 1)
58	- `--seq-length N`: Sequence length (default: 4096)
59	- `--num-gpus N`: Total GPU count (default: 8)
60	- `--recompute-granularity {full,selective}`: Enable activation checkpointing
61	- `--num-layers-in-first-pipeline-stage N`: Number of layers in the first pipeline stage (use when model layers cannot be evenly divided by `--pp`)
62	- `--num-layers-in-last-pipeline-stage N`: Number of layers in the last pipeline stage (use when model layers cannot be evenly divided by `--pp`)
63	- `--verbose`: Show detailed model breakdown
64	- `--json`: Output as JSON
65
66	Examples:
67
68	```bash
69	# Basic estimation
70	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 --num-gpus 64
71
72	# With memory optimization
73	python scripts/estimate_from_hf.py Qwen/Qwen3-235B-A22B \
74	--tp 8 --pp 4 --ep 4 \
75	--recompute-granularity full \
76	--recompute-method uniform \
77	--num-gpus 128
78
79	# Verbose output
80	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
81	--tp 4 --pp 4 --ep 8 --verbose --num-layers-in-last-pipeline-stage 16
82
83	# JSON output for automation
84	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
85	--tp 4 --pp 4 --ep 8 --json --num-layers-in-last-pipeline-stage 16 > result.json
86	```
87
88
89
90	## Common Workflows
91
92	### Find Optimal Parallelism for a Model
93
94	```bash
95	# Start with model path
96	MODEL="deepseek-ai/DeepSeek-V3"
97	GPUS=128
98
99	# Test different strategies
100	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 4 --ep 8 --num-gpus $GPUS --num-layers-in-last-pipeline-stage 16
101	python scripts/estimate_from_hf.py $MODEL --tp 8 --pp 2 --ep 8 --num-gpus $GPUS --num-layers-in-last-pipeline-stage 31
102
103
104	# Choose strategy that fits GPU memory with best efficiency
105	```
106
107	### Optimize for Memory Efficiency
108
109	Progressive memory reduction:
110
111	```bash
112	# 1. Baseline
113	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 2 --num-gpus 16
114
115	# 2. Add recomputation
116	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 2 --num-gpus 16 \
117	--recompute-granularity full
118
119	# 3. Increase expert parallelism (MoE only)
120	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 2 --ep 4 --num-gpus 16 \
121	--recompute-granularity full
122
123	# 4. Increase pipeline parallelism
124	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 4 --ep 4 --num-gpus 16 \
125	--recompute-granularity full
126
127	# 5. Last resort: reduce batch size
128	python scripts/estimate_from_hf.py $MODEL --tp 4 --pp 4 --ep 4 --num-gpus 16 \
129	--recompute-granularity full --micro-batch-size 1
130	```
131
132	### Check if Model Fits Available GPUs
133
134	```bash
135	# Check if DeepSeek-V3 fits in 128x A100 80GB
136	python scripts/estimate_from_hf.py deepseek-ai/DeepSeek-V3 \
137	--tp 4 --pp 4 --ep 8 --num-gpus 128 --num-layers-in-last-

People who liked this also viewed

6 picks

Type