diff --git a/README.md b/README.md
index acd95e6..d4510de 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,16 @@
# SimuMax: a static analytical model for LLM distributed training
-
+* [Latest News](#latest-news)
* [Introduction](#introduction)
* [Getting Started](#Installation)
* [Unitest](#Unitest)
* [TODO](#todo)
* [Acknowledgements](#acknowledgements)
* [Community](#Community)
+## Latest News
+- 📣 NEW! [App](./develop/app/README.md) - Add a user application to the SimuMax.
+- 📣 NEW! [Strategy Search](./examples/search/llm_search.py) - Support strategy search.
+- 📣 NEW! [System Config Pipeline](./simu_tools/efficency_test/README.md) - Add a Pipeline to generate system file include computing and communication efficiency.
+
## Introduction
SimuMax is a distributed training simulator designed for large-scale language model (LLM) workloads. It leverages a static analytical model to simulate and analyze both performance and memory usage, providing detailed insights into training efficiency without running the actual training process. Based on these analyses, SimuMax helps users explore potential ways to maximize computational efficiency.
@@ -35,6 +40,8 @@ It's appropriate to address various use-cases:
- [x] Megatron Compatibility: Simplified model migration pipeline
- [x] Finer-grained selective recompute
- [x] Efficiency measurement across shapes/layouts
+- [x] Efficiency measurement pipeline for both computing and commucation
+- [x] Strategy Search
### Benchmarks
@@ -42,8 +49,8 @@ Performance of some models on a single node. Llama3-70B was trimmed to 12 layers
Details can be found in [FULL_RESULTS](docs/FULL_RESULTS.md)
-
#### A100-Pcie
+Please refer to [./simu_tools/megatron_scripts/README.md](./simu_tools/megatron_scripts/README.md) for the benchmark script.
@@ -66,6 +73,8 @@ pip install -v -e .
## Example
+We recommend using our app program to get started with SimuMax with zero cost. Please refer to the [README](./app/README.md) for more details.
+
Please refer to the [tutorial](./docs/tutorial.md) for more details.
```bash
@@ -73,22 +82,44 @@ cd ./examples
python perf_llama3_8b_tp1_pp2.py
# The results are stored in the llama3_8b_a100_pcie_bf16 directory
```
-The output is as follows:
+The output is as follows, the output details are saved in the ./llama3_8b_a100_pcie_bf16 directory:
```
--------------SIMUMAX SUMMARY TP=1,EP=1,PP=2 -------------
-- parallelism = seq4096.mbs1.mbc8.gbs32 tp1.ep1.pp2.dp4.etp1.edp4, world_size:8
+-------------SIMUMAX SUMMARY llama3_8b(8.03B) TP=1,EP=1,PP=2 -------------
+- parallelism = layer32.dense0.seq4096.mbs1.mbc8.gbs32 tp1.ep1.pp2.dp4.etp1.edp4, world_size:8
- recompute = No Recompute
- dtype = bf16, grad_reduce = fp32
- system = a100_pcie_bf16
-- model = dense
-- mfu = 0.49
-- TFLOPS = 151.59 (tflops=843.3426 T, duration=5.5632 s)
-- TGS_per_gpu = 2945.052828675417
-- peak_alloc_mem = {'first_stage': '50.7760 GB', 'last_stage': '45.1637 GB'}
+- model_type = dense
+· mfu = 0.48
+· TFLOPS = 151.01T (tflops=843.3426 T, duration=5.5848 s)
+· TFLOPS_PER_TOKEN = 0.01T, duration=5.5848 s
+· peak_alloc_mem = {'first_stage': '50.8854 GB', 'last_stage': '45.1637 GB'}
+- peak_alloc_mem_with_reserved = {'first_stage': '54.1334 GB', 'last_stage': '48.0465 GB'}
+- TGS_per_gpu = 2933.6554333899194
+- net = pp_net=intra_node_pcie_8x, tp_net=intra_node_pcie_2x, dp_net=intra_node_pcie_4x, ep_net=intra_node_pcie_2x, etp_net=intra_node_pcie_2x
+------------------------------------------
```
-# Unitest
-after clone the repo, plz use "git config core.hooksPath git_hooks" to set commit hook, which will check the perf result in some config is still the same before commit automatically.
+## Strategy search example
+We provide a set of strategy search scripts that run on the [llm_search.py](./examples/search/llm_search.py) model, which can be used to search for the optimal strategy for your model.
+```
+cd examples/llm_search
+python llm_search.py
+```
+
+The example strategy search output results are saved in the directory ./search_deepseek_v3_a100_pcie_bf16 and search_deepseek_v2_a100_pcie_bf16.
+
+
+## Generate system file
+
+If you want to use SimuMax on new equipment, please refer to [./simu_tools/efficency_test/README.md](./simu_tools/efficency_test/README.md) to generate system files.
+
+## Benchmark
+We provide a set of benchmark testing scripts that run on NVIDIA GPUs, which can be used for SimuMax simulation calibration. These scripts enable actual measurements for llama and deepseek models. The benchmark results will be stored in a designated directory. For detailed benchmark procedures and interpretation of results, please refer to [./simu_tools/megatron_scripts/README.md](./simu_tools/megatron_scripts/README.md).
+
+## Notes
+- Currently, all Linear models are forced to perform gradient accumulation fusion.
+
# TODO
SimuMax is in active development, may contain bugs or incomplete features. Contributions are welcome!
@@ -97,7 +128,6 @@ There are features to be added. Several significants are:
- More pipeline scheduler
- Overlap between computation and communication
- Offloading
-- Strategy search
- More accurate memory-bound operator simulation
@@ -120,9 +150,4 @@ If you're passionate about:
Large-scale models for MoE, Reinforcement Learning, Multi-Modal
GPU/GPU-Cluster Training/Inference performance optimization
-Feel free to reach out to xuerong.huang@mthreads.com.
-
-## Star History
-## Star History
-
-[](https://www.star-history.com/#MooreThreads/SimuMax&type=date&legend=top-left)
+Feel free to reach out to 1354789084@qq.com.
\ No newline at end of file
diff --git a/app/README.md b/app/README.md
new file mode 100644
index 0000000..ba526ae
--- /dev/null
+++ b/app/README.md
@@ -0,0 +1,7 @@
+# Get started
+you can run the app with the following command:
+```shell
+cd app
+bash install.sh
+streamlit run streamlit_app.py
+```
\ No newline at end of file
diff --git a/app/install.sh b/app/install.sh
new file mode 100644
index 0000000..7822c3e
--- /dev/null
+++ b/app/install.sh
@@ -0,0 +1,7 @@
+# pip install PyQt5 -i https://pypi.tuna.tsinghua.edu.cn/simple
+# sudo apt-get update
+# sudo apt-get install libxcb-xinerama0 libxcb-icccm4 libxcb-image0 libxcb-keysyms1 libxcb-randr0 libxcb-render-util0 libxcb-shape0 libxcb-sync1 libxcb-xfixes0 libxcb-xkb1 libxkbcommon-x11-0
+
+
+# pip install flask -i https://pypi.tuna.tsinghua.edu.cn/simple
+pip install streamlit -i https://pypi.tuna.tsinghua.edu.cn/simple
\ No newline at end of file
diff --git a/app/streamlit_app.py b/app/streamlit_app.py
new file mode 100644
index 0000000..32c1062
--- /dev/null
+++ b/app/streamlit_app.py
@@ -0,0 +1,863 @@
+from functools import partial
+import streamlit as st
+import re
+import os
+import copy
+os.environ['SIMU_CHECK'] = '1'
+import json
+import zipfile
+from datetime import datetime
+from io import BytesIO
+from simumax.core.config import ParameterExtractor
+from simumax.core.utils import HumanReadableSize
+from simumax.core.perf_llm import PerfLLM
+from simumax.core.config import ModelConfig, StrategyConfig, SystemConfig
+
+# 设置页面配置
+st.set_page_config(
+ page_title="SimuMax分析工具",
+ page_icon="🚀",
+ layout="wide",
+ initial_sidebar_state="expanded"
+)
+perf_model = PerfLLM()
+
+class ParameterAnalyzer:
+ def __init__(self):
+ # 可选的配置预设
+ self.config_presets = {
+ 'small': {
+ 'seq_len': 512,
+ 'micro_batch_size': 8,
+ 'micro_batch_num': 2,
+ 'global_batch_size': 16,
+ 'tp_size': 2,
+ 'ep_size': 1,
+ 'pp_size': 2,
+ 'dp_size': 2,
+ 'world_size': 4
+ },
+ 'medium': {
+ 'seq_len': 1024,
+ 'micro_batch_size': 16,
+ 'micro_batch_num': 4,
+ 'global_batch_size': 64,
+ 'tp_size': 4,
+ 'ep_size': 1,
+ 'pp_size': 2,
+ 'dp_size': 4,
+ 'world_size': 32
+ },
+ 'large': {
+ 'seq_len': 2048,
+ 'micro_batch_size': 32,
+ 'micro_batch_num': 8,
+ 'global_batch_size': 256,
+ 'tp_size': 8,
+ 'ep_size': 1,
+ 'pp_size': 4,
+ 'dp_size': 8,
+ 'world_size': 256
+ }
+ }
+
+ # 硬件配置选项
+ from simumax.utils import RELEASE_MODELS,RELEASE_SYSTEM
+ self.simumax_hardware_options = {
+ 'A100-80GB-PCIE': RELEASE_SYSTEM['a100_pcie'],
+ }
+ # 模型规模选项
+ self.simumax_model_options = RELEASE_MODELS
+ self.simumax_model_options = {
+ 'deepseek_v2': RELEASE_MODELS['deepseekv2'],
+ 'deepseek_v3': RELEASE_MODELS['deepseekv3'],
+ 'llama3_8b': RELEASE_MODELS['llama3-8b'],
+ 'llama3_70b': RELEASE_MODELS['llama3-70b'],
+ 'llama3_405b': RELEASE_MODELS['llama3-405b_padding_128'],
+ }
+
+ def analyze_parameters(self, params, hardware_name, model_name):
+ """分析参数并返回结果"""
+ hardware = self.hardware_options[hardware_name]
+ model = self.model_options[model_name]
+
+ # 参数验证
+ warnings = []
+ recommendations = []
+
+ # 计算理论值
+ calculated_world_size = params['tp_size'] * params['pp_size'] * params['dp_size']
+ actual_global_batch_size = params['micro_batch_size'] * params['micro_batch_num'] * params['dp_size']
+
+ # 内存估算 (简化模型)
+ activation_memory = (params['seq_len'] * params['micro_batch_size'] *
+ params['tp_size'] * 4 * 12) / (1024 ** 3) # GB
+
+ model_memory = (model['params'] * 4) / (1024 ** 3) # 假设 FP32, GB
+
+ total_memory_estimate = activation_memory + model_memory
+
+ # 检查内存是否足够
+ if total_memory_estimate > hardware['memory'] * 0.9: # 留10%余量
+ warnings.append(f"内存估计 {total_memory_estimate:.2f}GB 超过硬件限制 {hardware['memory']}GB")
+ recommendations.append("考虑减小批次大小或使用模型并行")
+
+ # 检查配置一致性
+ if params['world_size'] != calculated_world_size:
+ warnings.append(f"world_size 配置不一致: 输入值 {params['world_size']}, 计算值 {calculated_world_size}")
+ recommendations.append(f"建议将 world_size 设置为 {calculated_world_size}")
+
+ if params['global_batch_size'] != actual_global_batch_size:
+ warnings.append(f"global_batch_size 配置不一致: 输入值 {params['global_batch_size']}, 计算值 {actual_global_batch_size}")
+ recommendations.append(f"建议将 global_batch_size 设置为 {actual_global_batch_size}")
+
+ # 性能估算
+ communication_overhead = (params['tp_size'] + params['pp_size']) * 0.05
+ efficiency_score = max(0, 100 - communication_overhead * 100)
+
+ # 吞吐量估算
+ estimated_throughput = (params['global_batch_size'] /
+ (1 + communication_overhead)) # tokens/step
+
+ return {
+ 'parameters': params,
+ 'analysis': {
+ 'calculated_world_size': calculated_world_size,
+ 'actual_global_batch_size': actual_global_batch_size,
+ 'memory_estimate_gb': round(total_memory_estimate, 2),
+ 'activation_memory_gb': round(activation_memory, 2),
+ 'model_memory_gb': round(model_memory, 2),
+ 'efficiency_score': round(efficiency_score, 1),
+ 'estimated_throughput': round(estimated_throughput, 2),
+ 'communication_overhead': round(communication_overhead, 2),
+ 'hardware_utilization': round((total_memory_estimate / hardware['memory']) * 100, 1),
+ 'warnings': warnings,
+ 'recommendations': recommendations,
+ 'is_config_valid': len(warnings) == 0
+ },
+ 'hardware_info': {'name': hardware_name, **hardware},
+ 'model_info': {'name': model_name, **model}
+ }
+
+def create_download_zip(perf_model:PerfLLM, mem_result, compute_result):
+ """创建下载文件"""
+ # 创建内存zip文件
+ memory_file = BytesIO()
+
+ with zipfile.ZipFile(memory_file, 'w') as zf:
+ self = perf_model
+ base_info = {}
+ base_info["arch"] = str(self.model_chunk_dict)
+ base_info["all_param"] = self.model_config.param_numel
+ base_info["act_param"] = self.model_config.activated_param_numel
+ # 添加文本报告
+ zf.writestr('model_arch.txt', base_info["arch"])
+ # 添加JSON配置
+ zf.writestr('base_info.json', json.dumps(base_info, indent=2, sort_keys=False, ensure_ascii=False))
+ zf.writestr('mem_result.json', str(mem_result))
+ zf.writestr('compute_result.json', str(compute_result))
+ zf.writestr('strategy_config.json', str(self.strategy))
+ zf.writestr('system_config.json', str(self.system))
+ zf.writestr('model_config.json', str(self.model_config))
+
+ memory_file.seek(0)
+ return memory_file
+
+def main():
+ # 初始化分析器
+ analyzer = ParameterAnalyzer()
+
+ # 页面标题
+ st.title("🚀 SimuMax分析工具")
+ st.markdown("分析和优化分布式训练配置参数")
+
+ # 侧边栏 - 快速配置
+ with st.sidebar:
+ st.header("⚡ 快速配置")
+
+ if 'selected_hardware' not in st.session_state:
+ st.session_state.selected_hardware = "A100-80GB-PCIE"
+ if 'selected_model' not in st.session_state:
+ st.session_state.selected_model = "deepseek_v3"
+
+ def update_hardware(main=False):
+ if main:
+ st.session_state.selected_hardware = st.session_state.main_hardware
+ else:
+ st.session_state.selected_hardware = st.session_state.side_hardware
+
+
+ def update_model(main=False):
+ if main:
+ st.session_state.selected_model = st.session_state.main_model
+ else:
+ st.session_state.selected_model = st.session_state.side_model
+
+
+ def update_paralism():
+ paralism = st.session_state.side_paralism
+ param_patterns = {
+ 'tp_size': (r'TP(\d+)', 1),
+ 'ep_size': (r'EP(\d+)', 1),
+ 'pp_size': (r'PP(\d+)', 1),
+ 'world_size': (r'GPU(\d+)', 8),
+ }
+ paralism_params = ParameterExtractor(param_patterns=param_patterns).extract_parameters(paralism)
+ print(paralism_params)
+ for key, value in paralism_params.items():
+ if key in st.session_state:
+ st.session_state[key] = value
+
+ # 硬件选择
+ side_hardware = st.selectbox(
+ "选择硬件配置",
+ list(analyzer.simumax_hardware_options.keys()),
+ index=list(analyzer.simumax_hardware_options.keys()).index(st.session_state.selected_hardware),
+ key="side_hardware",
+ on_change = partial(update_hardware, main=False),
+ )
+
+ # 模型选择
+ side_model = st.selectbox(
+ "选择模型规模",
+ list(analyzer.simumax_model_options.keys()),
+ key="side_model",
+ index=list(analyzer.simumax_model_options.keys()).index(st.session_state.selected_model),
+ on_change = partial(update_model, main=False),
+ )
+ init_paralisms = ['TP1+PP1+GPU8', 'TP1+PP2+GPU8', 'TP2+PP1+GPU8', 'EP4+PP2+GPU8', 'EP8+PP1+GPU8']
+ side_paralism = st.selectbox(
+ "选择并行方式",
+ init_paralisms,
+ key="side_paralism",
+ index=init_paralisms.index('EP8+PP1+GPU8'),
+ on_change = update_paralism
+ )
+
+ st.markdown("---")
+
+
+ # 硬件选择
+ main_hardware = st.selectbox(
+ "选择硬件配置",
+ # list(analyzer.hardware_options.keys())
+ list(analyzer.simumax_hardware_options.keys()),
+ index=list(analyzer.simumax_hardware_options.keys()).index(st.session_state.selected_hardware),
+ )
+
+ # 模型选择
+ main_model = st.selectbox(
+ "选择模型规模",
+ # list(analyzer.model_options.keys())
+ list(analyzer.simumax_model_options.keys()),
+ index=list(analyzer.simumax_model_options.keys()).index(st.session_state.selected_model),
+ )
+ perf_model.configure(
+ strategy_config=StrategyConfig.init_from_format_strings("gbs8"),
+ model_config=ModelConfig.init_from_config_file(analyzer.simumax_model_options[main_model]),
+ system_config=SystemConfig.init_from_config_file(analyzer.simumax_hardware_options[main_hardware])
+ )
+ st.success(f"✅ 已选择: {main_model}/{main_hardware}")
+
+ with st.expander("📋 模型详细信息", expanded=True):
+ detail_col1, detail_col2, detail_col3 = st.columns(3)
+ model_info = perf_model.model_config
+ with detail_col1:
+ st.write(f"**模型类型:** {model_info.model_type}")
+ st.write(f"**模型名称:** {model_info.model_name}")
+ st.write(f"**注意力类型:** {model_info.attention_type}")
+ st.write(f"**隐藏层大小:** {model_info.hidden_size}")
+ st.write(f"**头数量:** {model_info.head_num}")
+ st.write(f"**KV头数量:** {model_info.kv_head_num}")
+
+ with detail_col2:
+ st.write(f"**头大小:** {model_info.head_size}")
+ st.write(f"**中间隐藏层大小:** {model_info.intermediate_size}")
+ st.write(f"**总层数:** {model_info.layer_num}")
+
+ if model_info.model_type == 'moe':
+ st.write(f"**稠密层数:** {model_info.dense_layers}")
+ st.write(f"**专家数量:** {model_info.expert_num}")
+ st.write(f"**TopK:** {model_info.topk}")
+ st.write(f"**MoE FFN隐藏层大小:** {model_info.moe_ffn_hidden_size}")
+ st.write(f"**MoE共享专家隐藏层大小:** {model_info.moe_shared_expert_intermediate_size}")
+
+ with detail_col3:
+ if model_info.attention_type == 'mla':
+ st.write(f"**V注意力头维度:** {model_info.v_head_dim}")
+ st.write(f"**QK注意力头维度:** {model_info.qk_head_dim}")
+ st.write(f"**Q LoRA秩:** {model_info.q_lora_rank}")
+ st.write(f"**KV LoRA秩:** {model_info.kv_lora_rank}")
+ st.write(f"**QK位置编码维度:** {model_info.qk_pos_emb_head_dim}")
+
+ st.write(f"**词表大小:** {model_info.vocab_size}")
+ st.write(f"**使用SwiGLU:** {'是' if model_info.use_swiglu else '否'}")
+
+ st.markdown("---")
+ st.markdown("### 操作")
+ st.markdown("""
+
+ """, unsafe_allow_html=True)
+
+ analyze_btn = st.button("🚀 开始评估配置", use_container_width=True)
+ # analyze_btn = st.button("🎯 评估配置", use_container_width=True)
+ st.markdown("---")
+
+ st.markdown("### 关于")
+ st.info("""
+ 本工具用于分析分布式训练参数配置，
+ 提供内存估算、性能评估和优化建议。
+ """)
+
+ # 主内容区域
+ col1, col2 = st.columns([1, 1])
+
+ with col1:
+ st.header("🖥️ 硬件参数配置")
+
+ col_hw1, col_hw2, col_hw3 = st.columns(3)
+
+ with col_hw1:
+ st.subheader("网络通信")
+ if 'high_intra_node' in perf_model.system.networks:
+ intra_network_bandwidth = st.number_input(
+ "机内网络带宽 (GB/s)",
+ min_value=0.1,
+ max_value=1000.0,
+ value=st.session_state.get('intra_network_bandwidth', float(perf_model.system.networks['high_intra_node'].bandwidth.gbps)),
+ step=0.1,
+ help="节点内网络通信带宽"
+ )
+ if 'inter_node' in perf_model.system.networks:
+ inter_network_bandwidth = st.number_input(
+ "机间网络带宽 (GB/s)",
+ min_value=0.1,
+ max_value=1000.0,
+ value=st.session_state.get('inter_network_bandwidth', float(perf_model.system.networks['inter_node'].bandwidth.gbps)),
+ step=0.1,
+ help="节点间网络通信带宽"
+ )
+
+ with col_hw2:
+ st.subheader("计算能力")
+ compute_performance = st.number_input(
+ "算力 (TFLOPS)",
+ min_value=1.0,
+ max_value=1000.0,
+ value=float(perf_model.system.accelerator.op['matmul'].tflops),
+ step=1.0,
+ help="单卡理论计算性能"
+ )
+
+ # with col_hw3:
+ # st.subheader("算子效率")
+ # st.write("指定shape下的计算效率")
+ # op_efficiency_attn = st.slider(
+ # "注意力算子效率 (%)",
+ # min_value=10,
+ # max_value=100,
+ # value=65,
+ # help="注意力计算在实际shape下的效率"
+ # )
+ # op_efficiency_mlp = st.slider(
+ # "MLP算子效率 (%)",
+ # min_value=10,
+ # max_value=100,
+ # value=75,
+ # help="MLP计算在实际shape下的效率"
+ # )
+
+ st.header("📊 参数配置")
+
+ # 参数输入表格
+ with st.container():
+ st.subheader("训练参数")
+ train_col1_1, train_col1_2 = st.columns(2)
+
+ with train_col1_1:
+ seq_len = st.number_input(
+ "序列长度 (seq_len)",
+ min_value=1,
+ value=st.session_state.get('seq_len', 4096),
+ key='seq_len'
+ )
+
+ micro_batch_size = st.number_input(
+ "微批次大小 (mbs)",
+ min_value=1,
+ value=st.session_state.get('micro_batch_size', 1),
+ key='micro_batch_size'
+ )
+
+ global_batch_size = st.number_input(
+ "全局批次大小 (gbs)",
+ min_value=1,
+ value=st.session_state.get('global_batch_size', 256),
+ key='global_batch_size'
+ )
+
+ dtype = st.selectbox(
+ "数据类型",
+ ["bf16", "fp8"]
+ )
+
+ with train_col1_2:
+ tp_size = st.number_input(
+ "TP大小",
+ min_value=1,
+ value=st.session_state.get('tp_size', 1),
+ key='tp_size'
+ )
+
+ ep_size = st.number_input(
+ "EP大小",
+ min_value=1,
+ value=st.session_state.get('ep_size', 8),
+ key='ep_size'
+ )
+
+ pp_size = st.number_input(
+ "PP大小",
+ min_value=1,
+ value=st.session_state.get('pp_size', 1),
+ key='pp_size'
+ )
+ with st.expander("PP层数高级选项"):
+ num_layers_in_first_pipeline_stage = st.number_input(
+ "首个 Pipeline Stage的层数",
+ min_value=-1,
+ value=st.session_state.get('num_layers_in_first_pipeline_stage', -1),
+ key='num_layers_in_first_pipeline_stage',
+ help="如果为-1，则表示使用默认值"
+ )
+ num_layers_in_last_pipeline_stage = st.number_input(
+ "最后一个 Pipeline Stage的层数",
+ min_value=-1,
+ value=st.session_state.get('num_layers_in_last_pipeline_stage', -1),
+ key='num_layers_in_last_pipeline_stage',
+ help="如果为-1，则表示使用默认值"
+ )
+ world_size = st.number_input(
+ "卡数",
+ min_value=1,
+ value=st.session_state.get('world_size', 8),
+ key='world_size'
+ )
+
+ if 'previous_model' not in st.session_state:
+ st.session_state.previous_model = main_model
+
+ if st.session_state.previous_model != main_model:
+ model_config = perf_model.model_config
+ # 将dataclass转换为字典并批量更新
+ config_dict = model_config.__dict__
+ for key, value in config_dict.items():
+ if not key.startswith('_'): # 跳过私有属性
+ st.session_state[key] = value
+ st.session_state.previous_model = main_model
+ st.rerun()
+
+ with st.expander("🔽 模型参数配置"):#↓
+ model_col1_1, model_col1_2, model_col1_3 = st.columns(3)
+ with model_col1_1:
+ # noraml config
+ st.markdown("##### 🎯常规参数")
+ layer_num = st.number_input(
+ "层数",
+ min_value=1,
+ value=st.session_state.get('layer_num', perf_model.model_config.layer_num),
+ key='layer_num'
+ )
+ hidden_size = st.number_input(
+ "hidden_size",
+ min_value=1,
+ value=st.session_state.get('hidden_size', perf_model.model_config.hidden_size),
+ key='hidden_size'
+ )
+ intermediate_size = st.number_input(
+ "intermediate_size",
+ min_value=1,
+ value=st.session_state.get('intermediate_size', perf_model.model_config.intermediate_size),
+ key='intermediate_size'
+ )
+ vocab_size = st.number_input(
+ "vocab_size",
+ min_value=1,
+ value=st.session_state.get('vocab_size', perf_model.model_config.vocab_size),
+ key='vocab_size'
+ )
+ with model_col1_2:
+ # attention related
+ st.markdown("##### 👁️Attention参数")
+ head_num = st.number_input(
+ "head_num",
+ min_value=1,
+ value=st.session_state.get('head_num', perf_model.model_config.head_num),
+ key='head_num'
+ )
+ kv_head_num = st.number_input(
+ "kv_head_num",
+ min_value=1,
+ value=st.session_state.get('kv_head_num', perf_model.model_config.kv_head_num),
+ key='kv_head_num'
+ )
+ head_size = st.number_input(
+ "head_size",
+ min_value=1,
+ value=st.session_state.get('head_size', perf_model.model_config.head_size),
+ key='head_size'
+ )
+ if perf_model.model_config.attention_type == 'mla':
+ qk_head_dim = st.number_input(
+ "qk_head_dim",
+ min_value=1,
+ value=st.session_state.get('qk_head_dim', perf_model.model_config.qk_head_dim),
+ key='qk_head_dim'
+ )
+ v_head_dim = st.number_input(
+ "v_head_dim",
+ min_value=1,
+ value=st.session_state.get('v_head_dim', perf_model.model_config.v_head_dim),
+ key='v_head_dim'
+ )
+ qk_pos_emb_head_dim = st.number_input(
+ "qk_pose_emb_head_dim",
+ min_value=1,
+ value=st.session_state.get('qk_pose_emb_head_dim', perf_model.model_config.qk_pos_emb_head_dim),
+ key='qk_pose_emb_head_dim'
+ )
+ q_lora_rank = st.number_input(
+ "q_lora_rank",
+ min_value=1,
+ value=st.session_state.get('q_lora_rank', perf_model.model_config.q_lora_rank),
+ key='q_lora_rank'
+ )
+ kv_lora_rank = st.number_input(
+ "kv_lora_rank",
+ min_value=1,
+ value=st.session_state.get('kv_lora_rank', perf_model.model_config.kv_lora_rank),
+ key='kv_lora_rank'
+ )
+ if perf_model.model_config.model_type == 'moe':
+ with model_col1_3:
+ # moe related
+ st.markdown("##### 🏗️Moe参数")
+ dense_layers = st.number_input(
+ "dense_layers",
+ min_value=1,
+ value=st.session_state.get('dense_layers', perf_model.model_config.dense_layers),
+ key='dense_layers'
+ )
+ expert_num = st.number_input(
+ "expert_num",
+ min_value=1,
+ value=st.session_state.get('expert_num', perf_model.model_config.expert_num),
+ key='expert_num'
+ )
+ topk = st.number_input(
+ "topk",
+ min_value=1,
+ value=st.session_state.get('topk', perf_model.model_config.topk),
+ key='topk'
+ )
+ moe_ffn_hidden_size = st.number_input(
+ "moe_ffn_hidden_size",
+ min_value=1,
+ value=st.session_state.get('moe_ffn_hidden_size', perf_model.model_config.moe_ffn_hidden_size),
+ key='moe_ffn_hidden_size'
+ )
+ moe_shared_expert_intermediate_size = st.number_input(
+ "moe_shared_expert_intermediate_size",
+ min_value=1,
+ value=st.session_state.get('moe_shared_expert_intermediate_size', perf_model.model_config.moe_shared_expert_intermediate_size),
+ key='moe_shared_expert_intermediate_size'
+ )
+
+ st.subheader("重计算参数")#🔄
+ col1_1, col1_2 = st.columns(2)
+ with col1_1:
+ recompute_granularity = st.selectbox(
+ "重计算粒度",
+ options=[None, "selective_recompute", "full_recompute"],
+ format_func=lambda x: "无" if x is None else x,
+ key='recompute_granularity'
+ )
+ recompute_layer_num = st.number_input(
+ "重计算层数",
+ min_value=0,
+ value=st.session_state.get('recompute_layer_num', 0),
+ key='recompute_layer_num'
+ )
+ if recompute_granularity == "selective_recompute":
+ with col1_2:
+ attn_recompute = st.checkbox(
+ "ATTENTION重计算",
+ value=st.session_state.get('attn_recompute', False),
+ key='attn_recompute'
+ )
+ mla_rms_recompute = st.checkbox(
+ "MLA RMS重计算",
+ value=st.session_state.get('mla_rms_recompute', False),
+ key='mla_rms_recompute'
+ )
+
+ mlp_recompute = st.checkbox(
+ "MLP重计算",
+ value=st.session_state.get('mlp_recompute', False),
+ key='mlp_recompute'
+ )
+
+ mlp_rms_recompute = st.checkbox(
+ "MLP RMS重计算",
+ value=st.session_state.get('mlp_rms_recompute', False),
+ key='mlp_rms_recompute'
+ )
+ else:
+ attn_recompute = False
+ mla_rms_recompute = False
+ mlp_recompute = False
+ mlp_rms_recompute = False
+
+ with col2:
+ st.header("📈 分析结果")
+
+ # 显示当前选择的配置
+ st.info(f"**硬件:** {main_hardware} | **模型:** {main_model} | **并行:TP{tp_size}+EP{ep_size}+PP{pp_size}+GPU{world_size}**")
+
+ # 当点击分析按钮时执行分析
+ if analyze_btn:
+ with st.spinner("正在分析配置..."):
+ try:
+ # 1. set model config, refer 模型参数配置
+ ## normal model config
+ perf_model.model_config.layer_num = layer_num
+ perf_model.model_config.hidden_size = hidden_size
+ perf_model.model_config.intermediate_size = intermediate_size
+ perf_model.model_config.vocab_size = vocab_size
+ perf_model.strategy.dispatch_probs = True
+
+ ## attention model config
+ perf_model.model_config.head_num = head_num
+ perf_model.model_config.kv_head_num = kv_head_num
+ perf_model.model_config.head_size = head_size
+ if perf_model.model_config.attention_type == 'mla':
+ perf_model.model_config.qk_head_dim = qk_head_dim
+ perf_model.model_config.v_head_dim = v_head_dim
+ perf_model.model_config.qk_pos_emb_head_dim = qk_pos_emb_head_dim
+ perf_model.model_config.q_lora_rank = q_lora_rank
+ perf_model.model_config.kv_lora_rank = kv_lora_rank
+ if perf_model.model_config.model_type == 'moe':
+ perf_model.model_config.dense_layers = dense_layers
+ perf_model.model_config.expert_num = expert_num
+ perf_model.model_config.topk = topk
+ perf_model.model_config.moe_ffn_hidden_size = moe_ffn_hidden_size
+ perf_model.model_config.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size
+
+ # 2. set bw and tflops
+ if 'high_intra_node' in perf_model.system.networks:
+ perf_model.system.networks['high_intra_node'].bandwidth.gbps = intra_network_bandwidth
+ if 'inter_node' in perf_model.system.networks:
+ perf_model.system.networks['inter_node'].bandwidth.gbps = inter_network_bandwidth
+ perf_model.system.accelerator.op['default'].tflops = compute_performance
+ perf_model.system.accelerator.op['matmul'].tflops = compute_performance
+ perf_model.system.accelerator.op['fp8_matmul'].tflops = compute_performance
+ perf_model.system.accelerator.op['sdp_fwd'].tflops = compute_performance
+ perf_model.system.accelerator.op['sdp_bwd'].tflops = compute_performance
+ perf_model.system.accelerator.op['group_matmul'].tflops = compute_performance
+ perf_model.system.accelerator.op['fp8_group_matmul'].tflops = compute_performance
+
+ perf_model.model_config.moe_pad_expert_input_to_capacity = True
+ # TODO(sherry): add op efficiency
+
+ # 3. set paralilel strategy
+ perf_model.strategy.seq_len = seq_len
+ perf_model.strategy.micro_batch_size = micro_batch_size
+ perf_model.strategy.tp_size = tp_size
+ perf_model.strategy.ep_size = ep_size
+ perf_model.strategy.pp_size = pp_size
+ perf_model.strategy.world_size = world_size
+
+ if num_layers_in_last_pipeline_stage != -1:
+ perf_model.strategy.num_layers_in_last_pipeline_stage = num_layers_in_last_pipeline_stage
+ if num_layers_in_first_pipeline_stage != -1:
+ perf_model.strategy.num_layers_in_first_pipeline_stage = num_layers_in_first_pipeline_stage
+ perf_model.strategy.reset_global_batch_size(global_batch_size)
+
+
+ # 4. set recompute strategy
+ perf_model.strategy.enable_recompute = True
+ if recompute_granularity == 'full_recompute':
+ perf_model.strategy.recompute_granularity = 'full_block'
+ elif recompute_granularity == 'selective_recompute':
+ perf_model.strategy.recompute_granularity = 'selective_recompute'
+ else:
+ perf_model.strategy.recompute_granularity = None
+ perf_model.strategy.recompute_layer_num = recompute_layer_num
+ perf_model.strategy.attn_recompute = attn_recompute
+ perf_model.strategy.mla_rms_recompute = mla_rms_recompute
+ perf_model.strategy.mlp_recompute = mlp_recompute
+ perf_model.strategy.mlp_rms_recompute = mlp_rms_recompute
+
+ # 5. set dtype
+ perf_model.strategy.dtype = 'bf16'
+ if dtype == "fp8":
+ perf_model.strategy.fp8 = True
+
+ # 6. run estimate
+ perf_model.run_estimate()
+ result = perf_model.analysis()
+ mem_results = perf_model.analysis_mem().data
+ cost_results = perf_model.analysis_cost().data
+
+ st.session_state.analysis_result = (result, mem_results, cost_results, perf_model.strategy)
+ except Exception as e:
+ print(f"评估报错:{e}")
+ st.session_state.warnings = f"评估报错:{e}"
+
+ # 警告信息
+ if 'warnings' in st.session_state:
+ st.subheader("⚠️ 警告信息")
+ st.error(st.session_state.warnings)
+ # 删除警告信息
+ del st.session_state.warnings
+ elif 'analysis_result' in st.session_state: # 显示分析结果
+ try:
+ result, mem_results, cost_results, strategy = st.session_state.analysis_result
+ peak_mem = max(perf_model.get_pp_stage_peak_mem(mem_results, "peak_mem", False).values())
+ peak_mem_with_reserved = max(perf_model.get_pp_stage_peak_mem(mem_results, "peak_mem_with_reserved", False).values())
+
+ has_missed_op_efficiency = len(perf_model.system.miss_efficiency) > 0
+ if has_missed_op_efficiency:
+ missed_op_efficiency = copy.deepcopy(perf_model.system.miss_efficiency)
+ perf_model.system.reset_record_info()
+
+ overflow_memory = peak_mem_with_reserved/2**30 > perf_model.system.accelerator.mem_gbs
+ if overflow_memory or has_missed_op_efficiency:
+ st.subheader("💡 提示/建议") #❗️
+ # st.warning(f"**峰值Reserved显存({HumanReadableSize(peak_mem_with_reserved)})超过系统显存限制({perf_model.system.accelerator.mem_gbs}GB), 建议增加卡数或调整并行策略、重计算策略**")
+ warn_idx = 1
+ if overflow_memory:
+ st.markdown(
+ f'⚠️ {warn_idx}. 峰值Reserved显存({HumanReadableSize(peak_mem_with_reserved)})超过系统显存限制({perf_model.system.accelerator.mem_gbs}GB), 建议增加卡数或调整并行策略、重计算策略
',
+ unsafe_allow_html=True
+ )
+ warn_idx += 1
+ if has_missed_op_efficiency:
+ st.markdown(f'⚠️ {warn_idx}. 下面的op shape计算效率缺失,可能影响评估准确度,建议通过op测试脚本补全缺失shape的计算效率
',
+ unsafe_allow_html=True)
+ st.write(missed_op_efficiency)
+ warn_idx += 1
+
+ strategy:StrategyConfig = strategy
+ # 关键指标
+ st.subheader("📊 关键指标")
+ metric_col1, metric_col2, metric_col3 = st.columns(3)
+
+ with metric_col1:
+ st.metric("计算的卡数", strategy.world_size)
+ st.metric("内存估计" \
+ "(Peak Alloc)", f"{HumanReadableSize(peak_mem)}")
+
+ with metric_col2:
+ st.metric("实际全局批次大小", strategy.global_batch_size)
+ st.metric("MFU", f"{cost_results['mfu_6nd_with_attn']*100:.2f}%")
+
+ with metric_col3:
+ st.metric("Token吞吐量(TGS)", f"{cost_results['throughput_per_accelerator']:.2f}")
+ st.metric("算力吞吐量(TFLOPS)", f"{cost_results['throughput per GPU (TFLOP/s/GPU)']:.2f}")
+
+ # 内存细分
+ st.subheader("💾 内存分析")
+ if perf_model.strategy.pp_size == 1:
+ stages = ['first_stage']
+ elif perf_model.strategy.pp_size == 2:
+ stages = ['first_stage', 'last_stage']
+ elif perf_model.strategy.pp_size > 2:
+ stages = ['first_stage', 'middle_stage', 'last_stage']
+ pp_stage_labels = {
+ 'first_stage': 'Pipeline并行第一阶段',
+ 'middle_stage': 'Pipeline并行中间阶段',
+ 'last_stage': 'Pipeline并行最后阶段'
+ }
+ for stage in stages:
+ if perf_model.strategy.pp_size > 1:
+ context = st.expander(f"🔽 {pp_stage_labels[stage]}", expanded=True)
+ else:
+ context = st.container()
+ with context:#📁
+ if perf_model.strategy.pp_size > 1:
+ st.markdown(f"##### {pp_stage_labels[stage]}")
+ mem_result = mem_results[stage]
+ else:
+ mem_result = mem_results
+ mem_col1, mem_col2, mem_col3 = st.columns(3)
+
+ with mem_col1:
+ st.metric("前向激活内存(单Batch)", f"{mem_result['fwd_activation_cache_per_micro_batch']}")
+ st.metric("1F1B峰值激活内存(单Batch)", f"{mem_result['peak_activation_mem_in_1F1B']}")
+
+ with mem_col2:
+ st.metric("模型内存", f"{mem_result['model_mem']}")
+ with st.expander("📊 模型内存细分"):
+ st.write(f"- MoE部分: {mem_result['model_mem_detail']['moe']}")
+ st.write(f"- Dense部分: {mem_result['model_mem_detail']['dense']}")
+ with mem_col3:
+ st.metric("总峰值Alloc显存", f"{mem_result['peak_mem']}")
+ st.metric("总峰值Reserved显存", f"{mem_result['peak_mem_with_reserved']}")
+ # 配置验证状态
+ st.subheader("✅ 配置验证")
+ st.success("配置验证通过 ✓")
+
+ # st.warning("配置验证未通过 ⚠")
+ # st.subheader("💡 优化建议")
+ # with mem_col2:
+ # st.metric("模型内存", f"{analysis['model_memory_gb']} GB")
+ # with mem_col3:
+ # st.metric("总内存", f"{analysis['memory_estimate_gb']} GB")
+
+
+ # 显示resutlt信息
+ st.subheader("📊 汇总信息")
+ st.write(result)
+ with st.expander("详细通信带宽"):
+ st.write(perf_model.system.real_comm_bw)
+ # 下载报告
+ st.subheader("📥 下载报告")
+ zip_buffer = create_download_zip(perf_model, mem_results, cost_results)
+ st.download_button(
+ label="下载分析报告 (ZIP)",
+ data=zip_buffer,
+ file_name=f"training_analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip",
+ mime="application/zip",
+ use_container_width=True
+ )
+ # del st.session_state.analysis_result
+ except Exception as e:
+ print(f"分析报错:{e}")
+ st.subheader("⚠️ 警告信息")
+ st.error(f"分析报错:{e}")
+ else:
+ st.info("请点击'开始评估配置'按钮开始分析")
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/configs/models/deepseek-16b.json b/configs/models/deepseek-16b.json
new file mode 100644
index 0000000..d47ada7
--- /dev/null
+++ b/configs/models/deepseek-16b.json
@@ -0,0 +1,22 @@
+{
+ "model_type":"moe",
+ "model_name":"deepseek_16b",
+ "attention_type": "mla",
+ "hidden_size": 2048,
+ "head_num": 16,
+ "kv_head_num": 16,
+ "head_size": 128,
+ "intermediate_size": 10944,
+ "moe_ffn_hidden_size": 1408,
+ "moe_shared_expert_intermediate_size": 2816,
+ "layer_num": 28,
+ "expert_num": 64,
+ "v_head_dim": 128,
+ "qk_head_dim": 128,
+ "qk_pos_emb_head_dim": 64,
+ "q_lora_rank": null,
+ "kv_lora_rank": 512,
+ "topk":6,
+ "vocab_size": 102400,
+ "use_swiglu": true
+}
\ No newline at end of file
diff --git a/configs/models/deepseek-1b.json b/configs/models/deepseek-1b.json
new file mode 100644
index 0000000..97a5065
--- /dev/null
+++ b/configs/models/deepseek-1b.json
@@ -0,0 +1,25 @@
+{
+ "model_type": "moe",
+ "model_name":"deepseek_1b",
+ "attention_type": "mla",
+ "hidden_size": 2048,
+ "head_num": 16,
+ "kv_head_num": 16,
+ "head_size": 128,
+ "intermediate_size": 10944,
+ "moe_ffn_hidden_size": 1408,
+ "moe_shared_expert_intermediate_size": 2816,
+ "layer_num": 8,
+ "expert_num": 64,
+ "v_head_dim": 128,
+ "qk_head_dim": 128,
+ "qk_pos_emb_head_dim": 64,
+ "q_lora_rank": 1536,
+ "kv_lora_rank": 512,
+ "topk":6,
+ "vocab_size": 102400,
+ "use_swiglu": true,
+ "dense_layers" : 1,
+ "moe_pad_expert_input_to_capacity":true,
+ "capacity": 1
+}
\ No newline at end of file
diff --git a/configs/models/kimi-1T.json b/configs/models/kimi-1T.json
new file mode 100644
index 0000000..46b1fa4
--- /dev/null
+++ b/configs/models/kimi-1T.json
@@ -0,0 +1,23 @@
+{
+ "model_type":"moe",
+ "model_name":"kimi",
+ "attention_type": "mla",
+ "hidden_size": 7168,
+ "head_num": 64,
+ "kv_head_num": 64,
+ "head_size": 128,
+ "intermediate_size": 18432,
+ "moe_ffn_hidden_size": 2048,
+ "moe_shared_expert_intermediate_size": 2048,
+ "layer_num": 61,
+ "dense_layers": 1,
+ "expert_num": 384,
+ "v_head_dim": 128,
+ "qk_head_dim": 128,
+ "qk_pos_emb_head_dim": 64,
+ "q_lora_rank": 1536,
+ "kv_lora_rank": 512,
+ "topk":8,
+ "vocab_size": 163840,
+ "use_swiglu": true
+}
\ No newline at end of file
diff --git a/configs/models/qwen3-32b.json b/configs/models/qwen3-32b.json
new file mode 100644
index 0000000..df7c8ed
--- /dev/null
+++ b/configs/models/qwen3-32b.json
@@ -0,0 +1,12 @@
+{
+ "model_name": "qwen3-32b",
+ "model_type":"dense",
+ "hidden_size": 5120,
+ "head_num": 64,
+ "kv_head_num": 8,
+ "head_size": 128,
+ "intermediate_size": 25600,
+ "layer_num": 64,
+ "vocab_size": 151936,
+ "use_swiglu": true
+}
\ No newline at end of file
diff --git a/configs/strategy/ep4_pp2_dp4_mbs1.json b/configs/strategy/ep4_pp2_dp4_mbs1.json
index 45b096c..a008a93 100644
--- a/configs/strategy/ep4_pp2_dp4_mbs1.json
+++ b/configs/strategy/ep4_pp2_dp4_mbs1.json
@@ -14,11 +14,9 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94
}
\ No newline at end of file
diff --git a/configs/strategy/ep4_pp2_dp4_mbs1_full_recompute.json b/configs/strategy/ep4_pp2_dp4_mbs1_full_recompute.json
index a7e5f94..36435ff 100644
--- a/configs/strategy/ep4_pp2_dp4_mbs1_full_recompute.json
+++ b/configs/strategy/ep4_pp2_dp4_mbs1_full_recompute.json
@@ -14,12 +14,10 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94,
"recompute_granularity": "full_block",
"recompute_layer_num": 2
diff --git a/configs/strategy/ep4_pp2_dp4_mbs1_selective_recompute.json b/configs/strategy/ep4_pp2_dp4_mbs1_selective_recompute.json
index 47f5b52..40de56e 100644
--- a/configs/strategy/ep4_pp2_dp4_mbs1_selective_recompute.json
+++ b/configs/strategy/ep4_pp2_dp4_mbs1_selective_recompute.json
@@ -14,12 +14,10 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94,
"recompute_granularity": "selective_recompute",
"recompute_layer_num": 2,
diff --git a/configs/strategy/ep8_pp1_dp8_mbs1.json b/configs/strategy/ep8_pp1_dp8_mbs1.json
index 08b21a1..f419d34 100644
--- a/configs/strategy/ep8_pp1_dp8_mbs1.json
+++ b/configs/strategy/ep8_pp1_dp8_mbs1.json
@@ -14,11 +14,9 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94
}
\ No newline at end of file
diff --git a/configs/strategy/tp1_pp1_dp8_mbs1.json b/configs/strategy/tp1_pp1_dp8_mbs1.json
new file mode 100644
index 0000000..4db2570
--- /dev/null
+++ b/configs/strategy/tp1_pp1_dp8_mbs1.json
@@ -0,0 +1,22 @@
+{
+ "seq_len": 4096,
+ "micro_batch_size": 1,
+ "micro_batch_num": 8,
+ "dtype": "bf16",
+ "world_size": 8,
+ "tp_size": 1,
+ "pp_size": 1,
+ "ep_size": 1,
+ "etp_size": 1,
+ "moe_dispatcher_policy": "all2all",
+ "enable_sequence_parallel": true,
+ "interleaving_size": 1,
+ "zero_state": 1,
+ "enable_dropout": false,
+ "use_fused_norm": true,
+ "use_math_sdp": false,
+ "use_flash_sdp": true,
+ "use_fp32_accum_grad": true,
+ "enable_recompute": true,
+ "mem_factor": 0.94
+}
\ No newline at end of file
diff --git a/configs/strategy/tp1_pp2_dp4_mbs1.json b/configs/strategy/tp1_pp2_dp4_mbs1.json
index e1e43f0..44431ea 100644
--- a/configs/strategy/tp1_pp2_dp4_mbs1.json
+++ b/configs/strategy/tp1_pp2_dp4_mbs1.json
@@ -14,11 +14,9 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94
}
\ No newline at end of file
diff --git a/configs/strategy/tp2_pp1_dp4_mbs1.json b/configs/strategy/tp2_pp1_dp4_mbs1.json
index a6a46b2..c0eac25 100644
--- a/configs/strategy/tp2_pp1_dp4_mbs1.json
+++ b/configs/strategy/tp2_pp1_dp4_mbs1.json
@@ -14,11 +14,9 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94
}
\ No newline at end of file
diff --git a/configs/strategy/tp2_pp1_dp4_mbs1_full_recompute.json b/configs/strategy/tp2_pp1_dp4_mbs1_full_recompute.json
index 912df8f..cb2e275 100644
--- a/configs/strategy/tp2_pp1_dp4_mbs1_full_recompute.json
+++ b/configs/strategy/tp2_pp1_dp4_mbs1_full_recompute.json
@@ -14,12 +14,10 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94,
"recompute_granularity": "full_block",
"recompute_layer_num": 12
diff --git a/configs/strategy/tp2_pp1_dp4_mbs1_selective_recompute.json b/configs/strategy/tp2_pp1_dp4_mbs1_selective_recompute.json
index fb0a9a3..83e87ed 100644
--- a/configs/strategy/tp2_pp1_dp4_mbs1_selective_recompute.json
+++ b/configs/strategy/tp2_pp1_dp4_mbs1_selective_recompute.json
@@ -14,12 +14,10 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94,
"recompute_granularity": "selective_recompute",
"recompute_layer_num": 12,
diff --git a/configs/strategy/tp4_pp1_dp2_mbs1.json b/configs/strategy/tp4_pp1_dp2_mbs1.json
index 966ff05..794e4eb 100644
--- a/configs/strategy/tp4_pp1_dp2_mbs1.json
+++ b/configs/strategy/tp4_pp1_dp2_mbs1.json
@@ -14,11 +14,9 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94
}
\ No newline at end of file
diff --git a/configs/strategy/tp8_pp1_dp1_mbs1.json b/configs/strategy/tp8_pp1_dp1_mbs1.json
index ab13010..ed69800 100644
--- a/configs/strategy/tp8_pp1_dp1_mbs1.json
+++ b/configs/strategy/tp8_pp1_dp1_mbs1.json
@@ -14,11 +14,9 @@
"zero_state": 1,
"enable_dropout": false,
"use_fused_norm": true,
- "no_sync": true,
"use_math_sdp": false,
"use_flash_sdp": true,
"use_fp32_accum_grad": true,
"enable_recompute": true,
- "skip_ckpt_micro_batch_num": 0,
"mem_factor": 0.94
}
\ No newline at end of file
diff --git a/configs/system/a100_pcie.json b/configs/system/a100_pcie.json
index b74e9c1..155c1dc 100644
--- a/configs/system/a100_pcie.json
+++ b/configs/system/a100_pcie.json
@@ -242,18 +242,18 @@
"tflops": 312,
"efficient_factor": 0.75,
"accurate_efficient_factor":{
- "ng=40, M=616, N=3072, K=5120, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6313438865579614,
- "ng=40, M=616, N=3072, K=5120, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6790978664070304,
- "ng=40, M=616, N=3072, K=5120, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.5196854178569805,
- "ng=40, M=616, N=5120, K=1536, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6293014330752734,
- "ng=40, M=616, N=5120, K=1536, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6581820454471095,
- "ng=40, M=616, N=5120, K=1536, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.5198101456129606,
- "ng=20, M=1232, N=3072, K=5120, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.7073733774506031,
- "ng=20, M=1232, N=3072, K=5120, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.703850799964959,
- "ng=20, M=1232, N=3072, K=5120, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.6036380796120573,
- "ng=20, M=1232, N=5120, K=1536, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.675450005377,
- "ng=20, M=1232, N=5120, K=1536, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6572401208370234,
- "ng=20, M=1232, N=5120, K=1536, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.6372032871212036
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6313438865579614,
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6790978664070304,
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.5196854178569805,
+ "ng=40, M=616, N=5120, K=1536, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6293014330752734,
+ "ng=40, M=616, N=5120, K=1536, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6581820454471095,
+ "ng=40, M=616, N=5120, K=1536, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.5198101456129606,
+ "ng=20, M=1232, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.7073733774506031,
+ "ng=20, M=1232, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.703850799964959,
+ "ng=20, M=1232, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.6036380796120573,
+ "ng=20, M=1232, N=5120, K=1536, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.675450005377,
+ "ng=20, M=1232, N=5120, K=1536, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6572401208370234,
+ "ng=20, M=1232, N=5120, K=1536, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.6372032871212036
}
},
"fp8_group_matmul" : {
diff --git a/docs/FULL_RESULTS.md b/docs/FULL_RESULTS.md
index 6331f70..5007498 100644
--- a/docs/FULL_RESULTS.md
+++ b/docs/FULL_RESULTS.md
@@ -3,7 +3,6 @@
# Benchmarks
Performance of some models on a single node. Llama3-70B was trimmed to 12 layers and DeepSeek-236B was trimmed to 4 layers.
-
## A100-Pcie
diff --git a/docs/strategy-zh.md b/docs/strategy-zh.md
new file mode 100644
index 0000000..abc3123
--- /dev/null
+++ b/docs/strategy-zh.md
@@ -0,0 +1,90 @@
+
+ English|
+ 中文版本
+
+
+# 介绍
+SimuMax依赖于三个核心输入文件：system，strategy, model。strategy文件定义了训练的并行策略如tp/pp/ep，集群卡数， batchsize， 重计算策略等。
+
+# 参数说明
+## 基础训练参数
+### seq_len
+序列长度（token数量）
+### micro_batch_size
+微批次大小（单词每次前向传播处理的样本数）
+### micro_batch_num
+梯度累积的微批次数量
+### dtype
+计算数据类型（bf16表示半精度浮点数），默认为bf16
+### fp8
+是否使用fp8混合精度训练，默认为false
+## 分布式策略
+### world_size
+总GPU数量（默认8）
+### tp_size
+张量并行大小（Tensor Parallelism），默认为1
+### pp_size
+流水线大小（Pipeline Parallelism）### 纵向切分层数，默认为1
+### ep_size
+专家大小（Expert Parallelism）### 用于MOE模型，默认为1
+### etp_size
+专家张量大小（Expert Tensor Parallelism），默认为1
+### moe_dispatcher_policy
+MOE模型的路由策略， 默认为"all2all"
+### enable_sequence_parallel
+是否启用序列并行，默认为true，当tp_size > 1时生效
+### num_layers_in_first_pipeline_stage & num_layers_in_last_pipeline_stage
+控制第一个和最后一个Pipeline Parallel stage包含的层数，默认为None
+### interleaving_size
+保留字段，暂未使用
+### zero_state
+ZeRO优化配置，目前只支持zero0和1，默认为1
+## 内存优化
+### grad_reduce_in_bf16
+梯度归约是否使用bf16，默认为false
+### use_accm_weight
+是否使用累加权重融合（减少临时变量）, 默认为true
+### cache_groupgemm_col_fp8_inputs
+是否缓存groupgemm的FP8输入，默认为false
+### offload_groupgemm_col_inputs
+是否卸载groupgemm的输入到CPU，默认为false
+
+
+## 重计算相关
+#### enable_recompute
+recompute全局开关，默认为true
+#### recompute_granularity
+recompute的粒度，可选为"full_block"和"selective_recompute"，默认为None
+#### recompute_layer_num
+recompute的层数，默认为0
+#### attn_recompute
+对attention模块进行重计算，默认为false
+#### mla_rms_recompute
+对mla的rmsnorm和q/k up-projection进行重计算，默认为false
+#### mlp_recompute
+对MLP和groupedgemm进行重计算，默认为false
+#### mlp_rms_recompute
+对rmsnorm+router+sharedExpert进行重计算，默认为false
+#### recompute_variance
+recompute checkpoint的最后一个module是否去掉冗余的前向计算，默认为false, 当recompute_granularity为"selective_recompute"时，建议设置为true以节省计算时间
+
+## 计算优化
+### attention_sparse_rati
+注意力稀疏比例（0.0为密集注意力），默认为0.0
+### use_flash_sdp
+使用FlashAttention加速
+### use_fused_*
+各种融合内核优化
+### enable_dropout
+是否启用Dropout正则化，默认为false
+
+
+## 网络策略
+### tp_net, pp_net, dp_net等
+各种并行维度的网络通信策略，默认为"auto"，根据集群规模和并行策略自动选择
+
+## 其他
+### dispatch_probs
+Megatron-LM相关参数（是否分发probs）， 0.14版本Megatron-LM需要设置为true
+### mem_factor
+内存使用系数（0.94表示留6%的余量），用于估算reserve_memory（=max_memory / mem_factor），默认为0.94。
\ No newline at end of file
diff --git a/docs/strategy.md b/docs/strategy.md
new file mode 100644
index 0000000..d85499f
--- /dev/null
+++ b/docs/strategy.md
@@ -0,0 +1,91 @@
+
+ English|
+ 中文版本
+
+
+
+# Introduction
+SimuMax relies on three core input files: system, strategy, and model. The strategy file defines the training parallel strategies such as TP/PP/EP, cluster GPU count, batch size, recomputation strategies, etc.
+
+# Parameter Description
+## Basic Training Parameters
+### seq_len
+Sequence length (number of tokens)
+### micro_batch_size
+Micro-batch size (number of samples processed per forward propagation pass)
+### micro_batch_num
+Number of micro-batches for gradient accumulation
+### dtype
+Computation data type (bf16 indicates half-precision floating-point), default is bf16
+### fp8
+Whether to use fp8 mixed precision training, default is false
+
+## Distributed Strategy
+### world_size
+Total number of GPUs (default is 8)
+### tp_size
+Tensor Parallelism size, default is 1
+### pp_size
+Pipeline Parallelism size - vertically splits model layers, default is 1
+### ep_size
+Expert Parallelism size - used for MOE models, default is 1
+### etp_size
+Expert Tensor Parallelism size, default is 1
+### moe_dispatcher_policy
+Routing strategy for MOE models, default is "all2all"
+### enable_sequence_parallel
+Whether to enable sequence parallelism, default is true, effective when tp_size > 1
+### num_layers_in_first_pipeline_stage & num_layers_in_last_pipeline_stage
+Controls the number of layers contained in the first and last Pipeline Parallel stages, default is None
+### interleaving_size
+Reserved field, currently not used
+### zero_state
+ZeRO optimization configuration, currently only supports zero0 and zero1, default is 1
+
+## Memory Optimization
+### grad_reduce_in_bf16
+Whether to use bf16 for gradient reduction, default is false
+### use_accm_weight
+Whether to use weight accumulation fusion (reduces temporary variables), default is true
+### cache_groupgemm_col_fp8_inputs
+Whether to cache FP8 inputs for groupgemm, default is false
+### offload_groupgemm_col_inputs
+Whether to offload groupgemm inputs to CPU, default is false
+
+## Recompute Related
+#### enable_recompute
+Global switch for recompute, default is true
+#### recompute_granularity
+Granularity of recompute, options are "full_block" and "selective_recompute", default is None
+#### recompute_layer_num
+Number of layers for recompute, default is 0
+#### attn_recompute
+Recompute for attention module, default is false
+#### mla_rms_recompute
+Recompute for mla's rmsnorm and q/k up-projection, default is false
+#### mlp_recompute
+Recompute for MLP and groupedgemm, default is false
+#### mlp_rms_recompute
+Recompute for rmsnorm+router+sharedExpert, default is false
+#### recompute_variance
+Whether to remove redundant forward computation for the last module in recompute checkpoint, default is `false`.
+When recompute_granularity is "`selective_recompute`", it is recommended to set this to `true` to save computation time.
+## Computation Optimization
+### attention_sparse_ratio
+Attention sparse ratio (0.0 indicates dense attention), default is 0.0
+### use_flash_sdp
+Use FlashAttention acceleration
+### use_fused_*
+Various fused kernel optimizations
+### enable_dropout
+Whether to enable Dropout regularization, default is false
+
+## Network Strategy
+### tp_net, pp_net, dp_net, etc.
+Network communication strategies for various parallel dimensions, default is "auto", automatically selected based on cluster scale and parallel strategy
+
+## Other
+### dispatch_probs
+Megatron-LM related parameter (whether to dispatch probs), Megatron-LM version 0.14 requires this to be set to true
+### mem_factor
+Memory usage coefficient (0.94 means reserving 6% margin), used to estimate reserve_memory (=max_memory / mem_factor), default is 0.94
\ No newline at end of file
diff --git a/docs/system-zh.md b/docs/system-zh.md
new file mode 100644
index 0000000..f7cf8f6
--- /dev/null
+++ b/docs/system-zh.md
@@ -0,0 +1,252 @@
+
+ English|
+ 中文版本
+
+
+
+# 介绍
+SimuMax依赖于三个核心输入文件：system，strategy, model。system文件定义了一个GPU集群系统的硬件性能参数，例如标称算力，访存带宽，机内/机间带宽，用于评估计算时间和网络通信时间。
+
+system文件主要包含3部分描述：**基本信息(系统名称和单机卡数)**, **accelerator(算力和访存描述)**，**networks(网络配置和带宽描述)**。
+
+一个基础模板为:
+
+```json
+system_template = {
+ "sys_name": "A100",
+ "num_per_node": 8,
+ "accelerator": {
+ "backend": "cuda",
+ "mem_gbs": 80,
+ "op" : {
+
+ },
+ "bandwidth": {
+
+ },
+ "mode": "roofline"
+ },
+ "FC8":true,
+}
+```
+
+## accelerator
+accelerator部分包含了显存、访存带宽、算力、各个算子的计算效率等。
+
+### backend
+后端描述，仅用于标识。
+
+
+### mem_gbs
+显存大小，单位为GB。
+
+### op
+该部分定义了各个算子使用的默认算力和不同shape下准确的计算效率。
+
+SimuMax的核心之一是实现了shape级别计算效率建模，这是性能准确建模的关键，因此，SimuMax支持用户自定义多个核心算子在不同shape下的计算效率描述, 并且定义了一套shape表达规则， 用户需要按照该规则来新增算子不同shape的计算效率。
+
+目前支持的算子列表和其shape表达规则为：
+
+
+|key|算子|格式|示例|备注|
+|---|---|---|---|---|
+|matmul|矩阵乘法| b={batch_size}, m={m}, k={k}, n={n}, layout={layout}, accumulate={accumulate}, out_dtype={out_dtype}|`b=1, m=4096, k=5120, n=1536, layout=TN, accumulate=False, out_dtype=bf16`|accumulate：是否进行梯度累加，反向对w求导时该项为True|
+|fp8_matmul|FP8矩阵乘法|同上|同上|同上|
+|sdp_fwd|SDP前向计算|batch={batchh_size}, seq_len={seq_len}, head_num={head_num}, kv_head_num={kv_head_num}, qk_head_dim={qk_head_dim}, v_head_dim={v_head_dim}, qkv_contiguous={qkv_contiguous}|`batch=1, seq_len=4096, head_num=128, kv_head_num=128, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 1.0729673001633662`| qkv_contiguous：输入qkv是否在内存上连续，该项影响计算性能，因此单独描述，A100上一般为连续输入|
+|sdp_bwd|SDP反向计算|同上|同上|同上|
+|group_matmul| MOE模型的分组matmul|ng={num_groups}, M={fwd_M}, N={fwd_N}, K={fwd_k}, dtype={dtype}, out_dtype={out_dtype}, main_grad_dtype={main_grad_dtype}, stage={stage}, grad={grad}, accumulate={accumulate}, use_split_accumulator=False, single_output={single_output}| fwd stage：
`ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6313438865579614`|1. fwd、bwd_grad_act、bwd_grad_w三个阶段的groupedgemm shape描述的M,N,K都等于fwd阶段的M,N,K，用stage来区分不同阶段
2. single_output只有fwd stage时为True
3. accumulate只有bwd_grad_w stage时为True
4. grad和use_split_accumulator只有在bwd_grad_w stage时为True|
+
+
+
+
+例如，对于NVIDIA A100，其各个算子使用的算力和不同shape的计算效率描述为:
+
+```json
+
+"op": {
+ "default" : {
+ "tflops": 312,
+ "efficient_factor": 0.75
+ },
+ "matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "b=1, m=4096, k=5120, n=1536, layout=TN, accumulate=False, out_dtype=bf16": 0.7876672065615554,
+ "b=1, m=4096, k=1536, n=5120, layout=NN, accumulate=False, out_dtype=bf16": 0.737505124681297
+ },
+ },
+ "fp8_matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {},
+ },
+ "sdp_fwd" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "batch=1, seq_len=4096, head_num=128, kv_head_num=128, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 1.0729673001633662,
+ "batch=1, seq_len=4096, head_num=64, kv_head_num=64, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 1.0544285429372056
+ },
+ },
+ "sdp_bwd" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "batch=1, seq_len=4096, head_num=128, kv_head_num=128, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 0.8018473732899901,
+ "batch=1, seq_len=4096, head_num=64, kv_head_num=64, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 0.7942592665301026
+ },
+ },
+ "group_matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6313438865579614,
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6790978664070304,
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.5196854178569805
+ },
+ },
+ "fp8_group_matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {},
+ },
+}
+```
+其中，`default`表示默认算力，不支持的算子类型使用该算力；每个算子下面，`tflops`表示标称算力，`efficient_factor`表示默认计算效率，`accurate_efficient_factor`表示各个算子在不同shape下的实际计算效率。
+
+
+### bandwidth
+访存带宽描述，包含各个访存类型的带宽。例如，对于NVIDIA A100，其访存带宽描述为:
+
+```json
+"bandwidth": {
+ "default" : {
+ "efficient_factor": 0.91,
+ "gbps": 1600,
+ "latency_us": 40
+ },
+ "permute_fwd":{
+ "efficient_factor": 0.1879,
+ "gbps": 1600,
+ "latency_us": 40
+ },
+ "permute_bwd":{
+ "efficient_factor": 0.1879,
+ "gbps": 1600,
+ "latency_us": 40
+ },
+ "ce":{
+ "efficient_factor": 0.808,
+ "gbps": 1600,
+ "latency_us": 40
+ }
+}
+```
+其中，default表示默认访存带宽及其效率；除了默认访存带宽，我们新增3个memory bound算子的算子微调效率，用户可以自定义：
+- permute_fwd表示permute前向的访存带宽及其效率
+- permute_bwd表示permute反向的访存带宽及其效率
+- ce表示cross entropy的访存带宽及其效率
+
+## networks
+### FC8
+是否为FC8互联。
+### intra_with_pcie
+机内是否是pcie互连。
+- intra_with_pcie=True，则表示机内是pcie互连，networks还需包含以下网络带宽配置
+```json
+"intra_node_pcie_8x": {
+},
+"intra_node_pcie_4x": {
+},
+"intra_node_pcie_2x": {
+},
+"inter_node": {
+}
+```
+- intra_with_pcie=False， 否则表示机内是nvlink高速互连， networks还需包含以下网络带宽配置。
+```json
+"low_intra_node": {
+},
+"high_intra_node": {
+},
+"inter_node": {
+}
+```
+
+### intra_node_pcie_8x/intra_node_pcie_4x/intra_node_pcie_2x/low_intra_node/high_intra_node/inter_node
+每一种网络带宽配置，包含以下参数：
+- processor_usage: unused, 保留字段
+- bandwidth: 网络带宽配置，包含以下参数
+ - efficient_factor: 网络带宽效率
+ - gbps: 网络带宽
+ - latency_us: 网络延迟
+- op: 网络带宽效率，包含以下参数
+ - all_reduce: all_reduce操作的网络带宽效率
+ - scale: 2，固定
+ - offset: -1， 固定
+ - efficient_factor， 可选
+ - latency_us，可选
+ - all_gather: all_gather操作的网络带宽效率
+ - scale: 1，固定
+ - offset: -1， 固定
+ - efficient_factor， 可选
+ - latency_us，可选
+ - reduce_scatter: reduce_scatter操作的网络带宽效率
+ - scale: 1，固定
+ - offset: -1， 固定
+ - efficient_factor， 可选
+ - latency_us，可选
+ - p2p: p2p操作的网络带宽效率
+ - scale: 1，固定
+ - offset: -1， 固定
+ - efficient_factor， 可选
+ - latency_us，可选
+ - all2all: all2all操作的网络带宽效率
+ - scale: 1，固定
+ - offset: -1， 固定
+ - efficient_factor， 可选
+ - latency_us，可选
+
+例如A100_PCIE相邻两卡通信带宽详细配置：
+```json
+"intra_node_pcie_2x": {
+ "processor_usage": 0.00,
+ "bandwidth": {
+ "efficient_factor": 0.5,
+ "gbps": 30,
+ "latency_us": 10
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1,
+ "efficient_factor": 0.6965,
+ "latency_us": 15.51
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6866,
+ "latency_us": 24.84
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6419,
+ "latency_us": 131.30
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6969,
+ "latency_us": 24.07
+ }
+ }
+
+ },
+```
\ No newline at end of file
diff --git a/docs/system.md b/docs/system.md
new file mode 100644
index 0000000..047c668
--- /dev/null
+++ b/docs/system.md
@@ -0,0 +1,253 @@
+
+ English|
+ 中文版本
+
+
+
+# Introduction
+SimuMax relies on three core input files: system, strategy, and model. The system file defines the hardware performance parameters of a GPU cluster system, such as nominal computing power, memory access bandwidth, intra-node/inter-node bandwidth, used for evaluating computation time and network communication time.
+
+The system file primarily consists of three parts: Basic Information (system name and number of GPUs per node), accelerator (computing power and memory access description), and networks (network configuration and bandwidth description).
+
+A basic template is:
+
+```json
+system_template = {
+ "sys_name": "A100",
+ "num_per_node": 8,
+ "accelerator": {
+ "backend": "cuda",
+ "mem_gbs": 80,
+ "op" : {
+
+ },
+ "bandwidth": {
+
+ },
+ "mode": "roofline"
+ },
+ "FC8":true,
+}
+```
+
+## accelerator
+The accelerator section includes GPU memory size, memory access bandwidth, computing power, and computational efficiency for various operators.
+
+### backend
+Backend description, used for identification only.
+
+
+### mem_gbs
+GPU memory size, unit is GB.
+
+### op
+This section defines the default computing power used by various operators and the accurate computational efficiency under different shapes.
+
+One of SimuMax's core features is implementing shape-level computational efficiency modeling, which is key to accurate performance modeling. Therefore, SimuMax supports user-defined descriptions of computational efficiency for multiple core operators under different shapes and defines a set of shape expression rules. Users need to follow these rules to add computational efficiency for different operator shapes.
+
+The currently supported operator list and their shape expression rules are:
+
+
+
+
+|key|Operator|Format|Example|Notes|
+|---|---|---|---|---|
+|matmul|Matrix Multiplication| b={batch_size}, m={m}, k={k}, n={n}, layout={layout}, accumulate={accumulate}, out_dtype={out_dtype}|`b=1, m=4096, k=5120, n=1536, layout=TN, accumulate=False, out_dtype=bf16`|`accumulate`: whether gradient accumulation is performed; True during backward pass for w gradient|
+|fp8_matmul| FP8 Matrix Multiplication|Same as above|Same as above|Same as above|
+|sdp_fwd|SDP Forward Computation|batch={batchh_size}, seq_len={seq_len}, head_num={head_num}, kv_head_num={kv_head_num}, qk_head_dim={qk_head_dim}, v_head_dim={v_head_dim}, qkv_contiguous={qkv_contiguous}|`batch=1, seq_len=4096, head_num=128, kv_head_num=128, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 1.0729673001633662`| `qkv_contiguous`: whether input qkv is contiguous in memory; this affects performance, so described separately; generally contiguous input on A100|
+|sdp_bwd|SDP Backward Computation|Same as above|Same as above|Same as above|
+|group_matmul| Grouped matmul for MOE models|ng={num_groups}, M={fwd_M}, N={fwd_N}, K={fwd_k}, dtype={dtype}, out_dtype={out_dtype}, main_grad_dtype={main_grad_dtype}, stage={stage}, grad={grad}, accumulate={accumulate}, use_split_accumulator=False, single_output={single_output}| fwd stage：
`ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6313438865579614`|1. The` M, N, K `shape descriptions for the fwd, bwd_grad_act, bwd_grad_w stages all equal the M, N, K of the fwd stage, differentiated by the stage parameter.
2. `single_output` is True only for the fwd stage.
3. `accumulate` is True only for the bwd_grad_w stage.
4. `grad` and `use_split_accumulator` are True only for the bwd_grad_w stage.|
+|fp8_group_matmul| Grouped matmul for MOE models|Same as above|Same as above|Same as above|
+
+
+
+For example, for NVIDIA A100, the computing power used by its various operators and descriptions of computational efficiency under different shapes are:
+
+```json
+
+"op": {
+ "default" : {
+ "tflops": 312,
+ "efficient_factor": 0.75
+ },
+ "matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "b=1, m=4096, k=5120, n=1536, layout=TN, accumulate=False, out_dtype=bf16": 0.7876672065615554,
+ "b=1, m=4096, k=1536, n=5120, layout=NN, accumulate=False, out_dtype=bf16": 0.737505124681297
+ },
+ },
+ "fp8_matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {},
+ },
+ "sdp_fwd" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "batch=1, seq_len=4096, head_num=128, kv_head_num=128, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 1.0729673001633662,
+ "batch=1, seq_len=4096, head_num=64, kv_head_num=64, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 1.0544285429372056
+ },
+ },
+ "sdp_bwd" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "batch=1, seq_len=4096, head_num=128, kv_head_num=128, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 0.8018473732899901,
+ "batch=1, seq_len=4096, head_num=64, kv_head_num=64, qk_head_dim=192, v_head_dim=128, qkv_contiguous=True": 0.7942592665301026
+ },
+ },
+ "group_matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True": 0.6313438865579614,
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False": 0.6790978664070304,
+ "ng=40, M=616, N=3072, K=5120, dtype=bf16, out_dtype=bf16, main_grad_dtype=fp32, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False": 0.5196854178569805
+ },
+ },
+ "fp8_group_matmul" : {
+ "tflops": 312,
+ "efficient_factor": 0.75,
+ "accurate_efficient_factor": {},
+ },
+}
+```
+Here, `default` represents the default computing power, which is used for unsupported operator types; under each operator, `tflops` represents the nominal computing power, `efficient_factor` represents the default computational efficiency, and `accurate_efficient_factor` also indicates the actual computational efficiency of each operator under different shapes.
+
+
+### bandwidth
+Memory access bandwidth description, including bandwidth for various memory access types. For example, for NVIDIA A100, its memory access bandwidth description is:
+```json
+"bandwidth": {
+ "default" : {
+ "efficient_factor": 0.91,
+ "gbps": 1600,
+ "latency_us": 40
+ },
+ "permute_fwd":{
+ "efficient_factor": 0.1879,
+ "gbps": 1600,
+ "latency_us": 40
+ },
+ "permute_bwd":{
+ "efficient_factor": 0.1879,
+ "gbps": 1600,
+ "latency_us": 40
+ },
+ "ce":{
+ "efficient_factor": 0.808,
+ "gbps": 1600,
+ "latency_us": 40
+ }
+}
+```
+Here, `default` represents the default memory access bandwidth and its efficiency. Besides the default, we add operator-specific fine-tuned efficiency for 3 memory-bound operators that users can customize:
+- `permute_fwd` represents memory access bandwidth and efficiency for the permute forward pass.
+- `permute_bwd` represents memory access bandwidth and efficiency for the permute backward pass.
+- `ce` represents memory access bandwidth and efficiency for cross entropy.
+
+## networks
+### FC8
+Whether it is FC8 (Fully Connected 8) interconnect.
+### intra_with_pcie
+Whether intra-node connection uses PCIe.
+- if intra_with_pcie=True, it means intra-node connection uses PCIe, and the networks section must also include the following network bandwidth configurations:
+```json
+"intra_node_pcie_8x": {
+},
+"intra_node_pcie_4x": {
+},
+"intra_node_pcie_2x": {
+},
+"inter_node": {
+}
+```
+- if intra_with_pcie=False, it means intra-node connection uses high-speed NVLink, and the networks section must also include the following network bandwidth configurations:
+```json
+"low_intra_node": {
+},
+"high_intra_node": {
+},
+"inter_node": {
+}
+```
+
+### intra_node_pcie_8x/intra_node_pcie_4x/intra_node_pcie_2x/low_intra_node/high_intra_node/inter_node
+Each network bandwidth configuration includes the following parameters:
+- processor_usage: unused, reserved field
+- bandwidth: Network bandwidth configuration, includes:
+ - efficient_factor: Network bandwidth efficiency
+ - gbps: Network bandwidth, GB/s
+ - latency_us: Network latency
+- op: Network bandwidth efficiency for specific operations, includes:
+ - all_reduce: Network bandwidth efficiency for all_reduce operation
+ - scale: 2, fixed
+ - offset: -1， fixed
+ - efficient_factor， optional
+ - latency_us，optional
+ - all_gather: Network bandwidth efficiency for all_gather operation
+ - scale: 1, fixed
+ - offset: -1， fixed
+ - efficient_factor， optional
+ - latency_us，optional
+ - reduce_scatter: Network bandwidth efficiency for reduce_scatter operation
+ - scale: 1, fixed
+ - offset: -1， fixed
+ - efficient_factor， optional
+ - latency_us，optional
+ - p2p: Network bandwidth efficiency for p2p (point-to-point) operation
+ - scale: 1, fixed
+ - offset: -1， fixed
+ - efficient_factor， optional
+ - latency_us，optional
+ - all2all: Network bandwidth efficiency for all2all operation
+ - scale: 1, fixed
+ - offset: -1， fixed
+ - efficient_factor， optional
+ - latency_us，optional
+
+For example, detailed configuration of continuous two-card communication bandwidth for A100_PCIE:
+```json
+"intra_node_pcie_2x": {
+ "processor_usage": 0.00,
+ "bandwidth": {
+ "efficient_factor": 0.5,
+ "gbps": 30,
+ "latency_us": 10
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1,
+ "efficient_factor": 0.6965,
+ "latency_us": 15.51
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6866,
+ "latency_us": 24.84
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6419,
+ "latency_us": 131.30
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6969,
+ "latency_us": 24.07
+ }
+ }
+
+ },
+```
\ No newline at end of file
diff --git a/docs/tutorial.md b/docs/tutorial.md
index 24902a9..f91cfc4 100644
--- a/docs/tutorial.md
+++ b/docs/tutorial.md
@@ -26,6 +26,8 @@ cost_result = perf_model.analysis_cost()
```
In the above example, `system_config_file`, `strategy_config_file`, and `model_config_file` are paths to your configuration files.
+Please refer to [./system.md](./system.md) and [./strategy.md](./strategy.md) for more details on parameter configuration.
+
The run_estimate method simulates the training process and estimates the performance.
The analysis_mem method analyzes the memory usage during the training process and returns a mem_result object. This object contains information about the memory usage of different parts of the model, such as the weight memory usage, gradient memory usage, and state memory usage.
@@ -131,5 +133,26 @@ python perf_llama3_70b_layer12_tp2_full_recompute.py # the result is saved in d
```
-
-
+## Strategy Search
+### search_best_parallel_strategy_with_recompute
+The `search_best_parallel_strategy_with_recompute` interface can search for the best parallel strategy with the given recompute search space. The full example is as follows:
+```python
+perf_model = PerfLLM()
+perf_model.configure(
+ strategy_config=StrategyConfig.init_from_config_file(strategy_config_file),
+ model_config=ModelConfig.init_from_config_file(model_config_file),
+ system_config=SystemConfig.init_from_config_file(system_config_file),
+)
+perf_model.search_best_parallel_strategy_with_recompute(
+ world_size=2048,
+ gmi_error=6, # 6G memory reserved
+ micro_batch_size=1,
+ global_batch_size= 2048*8,
+ all_search_result=all_search_result,
+ tp_search_list=[1],
+ ep_search_list=[8, 16, 32, 64],
+ recompute_search_type=['no_recompute', 'full_block', 'selective_recompute'],
+ use_reserved_memory=False,
+ dump_path=f"search_{perf_model.model_config.model_name}_{perf_model.system.sys_name}"
+ )
+```
\ No newline at end of file
diff --git a/examples/search/llm_search.py b/examples/search/llm_search.py
new file mode 100644
index 0000000..3b87b38
--- /dev/null
+++ b/examples/search/llm_search.py
@@ -0,0 +1,53 @@
+import os
+# os.environ['ENABLE_SIMU_GRAPH'] = "1"
+from argparse import ArgumentParser
+from simumax.core.config import ModelConfig, StrategyConfig, SystemConfig
+from simumax.core.perf_llm import PerfLLM
+from simumax.utils import get_simu_model_config, get_simu_strategy_config, get_simu_system_config
+from simumax.core.graph import SimuONNXGraphBuilder, visualize_with_graphviz
+from simumax.core.tensor import FakeTensor
+
+def set_megatron_config(perf_model:PerfLLM):
+ perf_model.model_config.moe_pad_expert_input_to_capacity = True
+ perf_model.model_config.capacity = 1
+ perf_model.model_config.padded_vocab_size = True
+ perf_model.model_config.make_vocab_size_divisible_by = 128
+
+ perf_model.strategy.dispatch_probs = True
+ # perf_model.strategy.fp8 = True
+
+def search(model, system):
+ # Use the interface of SimuMax to configure the model_config_path, strategy_config_path and system_config_path, you can alse pase the path directly, sush as:
+ # system_config_file = '../configs/system/a100_pcie.json'
+ strategy_config_file = get_simu_strategy_config('tp1_pp2_dp4_mbs1') # default is tp1_pp2_dp4_mbs1
+ model_config_file = get_simu_model_config(model)
+ system_config_file = get_simu_system_config(system)
+
+
+ perf_model = PerfLLM()
+ perf_model.configure(
+ strategy_config=StrategyConfig.init_from_config_file(strategy_config_file),
+ model_config=ModelConfig.init_from_config_file(model_config_file),
+ system_config=SystemConfig.init_from_config_file(system_config_file),
+ )
+
+
+ set_megatron_config(perf_model) # Optinionally, just for align the configuration with the benchmark
+
+ all_search_result = {}
+ perf_model.search_best_parallel_strategy_with_recompute(
+ world_size=2048,
+ gmi_error=6, # 6G memory reserved
+ micro_batch_size=1,
+ global_batch_size= 2048*8,
+ all_search_result=all_search_result,
+ tp_search_list=[1],
+ ep_search_list=[8, 16, 32, 64],
+ recompute_search_type=['no_recompute', 'full_block', 'selective_recompute'],
+ use_reserved_memory=False,
+ dump_path=f"search_{perf_model.model_config.model_name}_{perf_model.system.sys_name}"
+ )
+
+if __name__ == "__main__":
+ search('deepseekv3', 'a100_pcie')
+ search('deepseekv2', 'a100_pcie')
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index c8d1917..1d2f6d3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,4 +4,5 @@ matplotlib
numpy
pandas
sympy
-tabulate
\ No newline at end of file
+tabulate
+graphviz
\ No newline at end of file
diff --git a/simu_tools/efficency_test/README-zh.md b/simu_tools/efficency_test/README-zh.md
new file mode 100644
index 0000000..cc9bbb7
--- /dev/null
+++ b/simu_tools/efficency_test/README-zh.md
@@ -0,0 +1,88 @@
+
+ English|
+ 中文版本
+
+
+# 自动生成system.json文件
+## 介绍
+SimuMax的性能评估准确性不仅得益于标准的模型定义、与训练框架一致的计算数据流，还得益于算子shape粒度的实测计算效率和逐通信原语的带宽效率+latency拟合。
+
+SimuMax的system.json文件定义了设备的标称算力、各种算子在不同shape下的计算效率、机内/机间带宽等，其各个属性描述请查阅[system.md](../../docs/system.md)。
+
+## 步骤
+生成符合SimuMax标准的system.json文件包含两个步骤：
+1. 测试算子shape级别的计算效率：首先批量运行给定的标准算子计算效率测试流程得到system.json初版
+2. 拟合机内/机间通信带宽：运行例如nccl_test通信测试工具，得到各个通信原语在一个通信量区间的通信时间，基于该通信量区间的实测时间拟合各个通信算子的机内/机间通信效率和latency。将拟合的通信效率和latency填入system.json中，完成。
+
+### 测试算子shape级别的计算效率
+首先批量运行给定的标准算子计算效率测试流程得到system.json初版。
+
+根据硬件实际配置，修改[run.sh](./run.sh)文件中的MAX_TFLOPS、SYS_NAME、PICE_INTRA_LINK等环境变量来指定设备的标称算力、系统名称、是否卡键为PCIE互联等。详细示例如下：
+```shell
+export MAX_TFLOPS=312 # 指定机器的标称算力
+export SYS_NAME="A100_PCIE" #指定system名称
+export PICE_INTRA_LINK=1 #指定机内卡间是否为PCIE互联
+export FC8_MODE=0 #指定机内卡间通信连接是否为FC8模式
+export PARAM_FILE="./run_params.json" # 测试超参数，指定测试模型的列表，mbs区间，seq_len区间等, 如果不指定，则使用默认超参。
+python test_gemm_efficency.py
+python test_grouped_gemm_efficency.py
+python test_fa_efficency.py
+python combine_efficency.py
+```
+其中run_params.json文件示例如下：
+```json
+{
+ "model_list":["deepseekv2",
+ "deepseekv3",
+ "deepseek-32b",
+ "deepseek-16b",
+ "deepseek-1b",
+ "llama3-8b",
+ "llama3-70b",
+ "qwen3-32b",
+ "kimi-1T",
+ "mixtral-8x7b"],
+ "mbs_list": [1, 2, 4],
+ "seq_len_list":[4096],
+ "tp_list": [1, 2, 4, 8],
+ "ep_list": [1, 2, 4, 8, 16, 64]
+}
+```
+运行：
+```shell
+bash run.sh
+```
+在运行目录下会生成{system_name}.json文件（例如实例中生成A100_PCIE.json文件）。每个模型的每种超参组合测试运行时间通常大约在1min左右。
+
+
+生成的system文件中，"networks"部分的机内带宽数值为默认的经验值，如果PICE_INTRA_LINK=0，则为A100_pcie的经验值：30GB/s的机内标称带宽及其精调效率+200GB/s的机间带宽；如果PICE_INTRA_LINK=1，则为A100-sxm的50GB/s机内标称带宽+200GB/s的机间带宽。
+
+注意：现在该测试框架基于简单的时间戳统计算子运行时间，因此测试结果可能存在误差，后续会考虑使用更精确的测试方法，例如 cuda events。
+
+#### 效率自动化测试框架支持的模型限制
+SimuMax目前可仿真的模型在./configs/models目录下，但不是所有模型的所有超参配置都可以基于现在这套框架进行自动化算子效率测试：
+- Moe模型：对于MOE模型，仅支持TP=1配置下的效率测试，这意味着对于mixtral这类可开TP(ETP)的模型。
+- Dense模型：全支持
+
+### 拟合机内/机间通信带宽
+如果设备为A100_PCIE，则上述步骤生成的system文件中，默认"networks"部分的机内带宽为我们已经拟合过的带宽，可以直接使用，如果多机进行perf，仍需要进一步对机间带宽精调。
+
+如果为其它设备，则需要使用例如[nccl_test](https://github.com/NVIDIA/nccl-tests)工具来测试reduce-scatter、all2all、allgather等通信算子在某一区间（建议1M->8GB）的通信时间，随后通过线性回归拟合一个bw*eff和latency的带宽模型，得到各个通信原语的通信效率值eff和latency。
+
+- nccl_test测试命令参考[nccl_test.sh](nccl_test.sh)：
+
+ ```shell
+ end=8G
+ echo "run all_reduce_perf"
+ ./build/all_reduce_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_all_reduce.txt
+ echo "run all_gather_perf"
+ ./build/all_gather_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_all_gather.txt
+ echo "run reduce_scatter_perf"
+ ./build/reduce_scatter_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_reduce_scatter.txt
+ echo "run alltoall_perf"
+ ./build/alltoall_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_alltoall_perf.txt
+ echo "run sendrecv_perf"
+ ./build/sendrecv_perf -n 1 -b 1M -e $end -f 2 -g 1 -t 1 -d bfloat16 > send_recv.txt
+ ```
+
+- 基于测试通信时间拟合通信效率和latency，请参考[nccl_fit.py](nccl_fit.py)。
\ No newline at end of file
diff --git a/simu_tools/efficency_test/README.md b/simu_tools/efficency_test/README.md
new file mode 100644
index 0000000..2618420
--- /dev/null
+++ b/simu_tools/efficency_test/README.md
@@ -0,0 +1,102 @@
+
+ English|
+ 中文版本
+
+
+# Automatic Generation of system.json File
+
+## Introduction
+The accuracy of SimuMax's performance evaluation benefits not only from the standard model definition and computation data flow consistent with training frameworks, but also from the measured computational efficiency at the operator-shape granularity and the bandwidth efficiency and latency fitting of each communication primitive.
+
+The system.json file in SimuMax defines the device's nominal computing power, computational efficiency of various operators under different shapes, intra-node/inter-node bandwidth, etc. For descriptions of each attribute, please refer to [system.md](../../docs/system.md).
+
+We provide a standard testing process that can generate system files that meet SimuMax's requirements.
+
+## Steps
+Generating a system.json file that complies with SimuMax standards consists of two steps:
+
+Test computational efficiency at operator-shape level: First, batch-run the standard operator computational efficiency testing process to obtain an initial version of system.json.
+
+Fit intra-node/inter-node communication bandwidth: Run tools such as nccl communication test tools to obtain communication times for each communication primitive within a communication volume range. Based on the measured times within this range, fit the intra-node/inter-node communication efficiency and latency for each communication operator. Fill the fitted communication efficiency and latency values into system.json to complete the process.
+
+### Step 1: Test Computational Efficiency at Operator-Shape Level
+First, batch-run the standard operator computational efficiency testing process to obtain an initial version of system.json.
+
+According to the actual hardware configuration, modify environment variables in [run.sh](./run.sh) such as MAX_TFLOPS, SYS_NAME, PICE_INTRA_LINK, etc., to specify the device's nominal computing power, system name, whether cards are connected via PCIE, etc. A detailed example is as follows:
+
+```shell
+export MAX_TFLOPS=312 # Specify the machine's nominal computing power
+export SYS_NAME="A100_PCIE" # Specify the system name
+export PICE_INTRA_LINK=1 # Specify whether intra-machine cards are connected via PCIE
+export FC8_MODE=0 # Specify whether intra-machine communication connections are in FC8 mode
+export PARAM_FILE="./run_params.json" # Test hyperparameters, specify the list of test models, mbs interval, seq_len interval, etc. If not specified, the default hyperparameters will be used.
+python test_gemm_efficency.py
+python test_grouped_gemm_efficency.py
+python test_fa_efficency.py
+python combine_efficency.py
+```
+An example of the run_params.json file is as follows:
+```json
+{
+ "model_list":["deepseekv2",
+ "deepseekv3",
+ "deepseek-32b",
+ "deepseek-16b",
+ "deepseek-1b",
+ "llama3-8b",
+ "llama3-70b",
+ "qwen3-32b",
+ "kimi-1T",
+ "mixtral-8x7b"],
+ "mbs_list": [1, 2, 4],
+ "seq_len_list":[4096],
+ "tp_list": [1, 2, 4, 8],
+ "ep_list": [1, 2, 4, 8, 16, 64]
+}
+```
+
+Run:
+```shell
+bash run.sh
+```
+A {system_name}.json file will be generated in the running directory (e.g., A100_PCIE.json in this example).
+
+In the generated system file, the intra-machine bandwidth values in the "networks" section are default empirical values.
+
+If `PICE_INTRA_LINK=0`, it corresponds to the empirical values for A100_pcie:
+- 30 GB/s nominal intra-machine bandwidth with finely-tuned efficiency
+- Plus 200 GB/s inter-machine bandwidth
+
+If `PICE_INTRA_LINK=1`, it corresponds to:
+- 50 GB/s nominal intra-machine bandwidth for A100-sxm
+- Plus 200 GB/s inter-machine bandwidth
+
+Note: The current testing framework calculates operator runtime based on simple timestamping, so the test results may contain errors. More precise testing methods, such as CUDA events, will be considered for future implementation.
+
+#### Model Limitations Supported by the Efficiency Automation Testing Framework
+The models that SimuMax can currently simulate are located in the ./configs/models directory, but not all hyperparameter configurations of all models can be tested for operator efficiency automatically using the current framework:
+- MoE Models: For MoE models, efficiency testing is only supported under TP=1 configurations, which means for models like Mixtral that allow TP (ETP).
+- Dense Models: Fully supported
+
+### Step 2: Fit intra-node/inter-node Communication Bandwidth
+If the device is A100_PCIE, then **in the system file generated through the above steps, the default intra-machine bandwidth in the "networks" section uses our pre-calibrated bandwidth values and can be used directly. However, if performing multi-machine performance simulation, further fine-tuning of inter-machine bandwidth is still required.**
+
+For other devices, tools such as [nccl_test](https://github.com/NVIDIA/nccl-tests) need to be used to test communication times for communication operators like reduce-scatter, all2all, allgather, etc., within a certain range (recommended: 1M to 8GB). Then, use linear regression to fit a bandwidth model with bw*eff and latency, obtaining the communication efficiency value eff and latency for each communication primitive.
+
+Reference for nccl_test commands: [nccl_test.sh](nccl_test.sh):
+
+```shell
+end=8G
+echo "run all_reduce_perf"
+./build/all_reduce_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_all_reduce.txt
+echo "run all_gather_perf"
+./build/all_gather_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_all_gather.txt
+echo "run reduce_scatter_perf"
+./build/reduce_scatter_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_reduce_scatter.txt
+echo "run alltoall_perf"
+./build/alltoall_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_alltoall_perf.txt
+echo "run sendrecv_perf"
+./build/sendrecv_perf -n 1 -b 1M -e $end -f 2 -g 1 -t 1 -d bfloat16 > send_recv.txt
+```
+
+For fitting communication efficiency and latency based on tested communication times, please refer to [nccl_fit.py](nccl_fit.py).
\ No newline at end of file
diff --git a/simu_tools/efficency_test/combine_efficency.py b/simu_tools/efficency_test/combine_efficency.py
new file mode 100644
index 0000000..05c72fc
--- /dev/null
+++ b/simu_tools/efficency_test/combine_efficency.py
@@ -0,0 +1,319 @@
+import os
+import json
+from simu_tools.efficency_test.utils import get_system_name
+
+system, device, MAX_TFLOPS = get_system_name()
+gemm_save_root = f'{system}_gemm_efficency/gemm_efficency.json'
+groupgemm_save_root = f'{system}_grouped_gemm_efficency/grouped_gemm_efficency.json'
+fa_save_root = f'{system}_fa_efficency/fa_efficency_test.json'
+
+gemm_ops = json.load(open(gemm_save_root, 'r'))
+groupgemm_ops = json.load(open(groupgemm_save_root, 'r'))
+fa_ops = json.load(open(fa_save_root, 'r'))
+
+sys_name = os.environ.get('SYS_NAME', 'A100_pcie')
+intra_with_pcie = int(os.environ.get('PICE_INTRA_LINK', '0'))
+FC8 = bool(int(os.environ.get('FC8_MODE', '0')))
+
+networks = {
+ "intra_with_pcie": False,
+ "low_intra_node": {
+ "processor_usage": 0.0,
+ "bandwidth": {
+ "efficient_factor": 1,
+ "gbps": 21,
+ "latency_us": 130
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1
+ }
+ }
+ },
+ "high_intra_node": {
+ "processor_usage": 0.0,
+ "bandwidth": {
+ "efficient_factor": 0.5,
+ "gbps": 50,
+ "latency_us": 35,
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1,
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1,
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1,
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0,
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1,
+ }
+ }
+ },
+ "inter_node": {
+ "processor_usage": 0.0,
+ "bandwidth": {
+ "efficient_factor": 1,
+ "gbps": 200,
+ "latency_us": 35
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1,
+ }
+ }
+ }
+ }
+pcie_networks = {
+ "intra_with_pcie":True,
+ "intra_node_pcie_8x": {
+ "processor_usage": 0.00,
+ "bandwidth": {
+ "efficient_factor": 0.5,
+ "gbps": 30,
+ "latency_us": 10
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1,
+ "efficient_factor": 0.4176,
+ "latency_us": 8.98
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.4497,
+ "latency_us": 4.83,
+ "dp_fixed_bw": {
+ "2": 4,
+ "4": 7.65
+ }
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.3688,
+ "latency_us": 15.64,
+ "dp_fixed_bw":{
+ "2": 3.9,
+ "4": 7
+ }
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.2754,
+ "latency_us": 30.04
+ }
+ }
+ },
+ "intra_node_pcie_4x": {
+ "processor_usage": 0.00,
+ "bandwidth": {
+ "efficient_factor": 0.5,
+ "gbps": 30,
+ "latency_us": 10
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1,
+ "efficient_factor": 0.6924,
+ "latency_us": 3.45
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6915,
+ "latency_us": 8.93
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6845,
+ "latency_us": 29.88
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6380,
+ "latency_us": 13.84
+ }
+ }
+
+ },
+ "intra_node_pcie_2x": {
+ "processor_usage": 0.00,
+ "bandwidth": {
+ "efficient_factor": 0.5,
+ "gbps": 30,
+ "latency_us": 10
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1,
+ "efficient_factor": 0.6965,
+ "latency_us": 15.51
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6866,
+ "latency_us": 24.84
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6419,
+ "latency_us": 131.30
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1,
+ "efficient_factor": 0.6969,
+ "latency_us": 24.07
+ }
+ }
+
+ },
+ "inter_node": {
+ "processor_usage": 0.00,
+ "bandwidth": {
+ "efficient_factor": 0.6,
+ "gbps": 200,
+ "latency_us": 130
+ },
+ "op": {
+ "all_reduce": {
+ "scale": 2,
+ "offset": -1
+ },
+ "all_gather": {
+ "scale": 1,
+ "offset": -1
+ },
+ "reduce_scatter": {
+ "scale": 1,
+ "offset": -1
+ },
+ "p2p": {
+ "scale": 1,
+ "offset": 0
+ },
+ "all2all": {
+ "scale": 1,
+ "offset": -1
+ }
+ }
+ }
+ }
+
+system_template = {
+ "sys_name": sys_name,
+ "num_per_node": 8,
+ "accelerator": {
+ "backend": "cuda",
+ "mem_gbs": 80,
+ "op" : {
+
+ },
+ "bandwidth": {
+ "default": {
+ "efficient_factor": 0.666,
+ "gbps": 1600,
+ "latency_us": 30
+ },
+ "permute_fwd":{
+ "efficient_factor": 0.46,
+ "gbps": 1600,
+ "latency_us": 30
+ },
+ "permute_bwd":{
+ "efficient_factor": 0.2175,
+ "gbps": 1600,
+ "latency_us": 30
+ },
+ "ce":{
+ "efficient_factor": 0.73,
+ "gbps": 1600,
+ "latency_us": 30
+ }
+ },
+ "mode": "roofline"
+ },
+ "FC8":FC8,
+}
+if intra_with_pcie:
+ system_template['networks'] = pcie_networks
+else:
+ system_template['networks'] = networks
+
+system_template["accelerator"]['op'].update(gemm_ops)
+system_template["accelerator"]['op'].update(groupgemm_ops)
+system_template["accelerator"]['op'].update(fa_ops)
+
+bf16_efficient_factor = [v['efficient_factor'] for k, v in system_template["accelerator"]['op'].items() if 'fp8' not in k and 'sdp' not in k]
+avg_efficient_factor = sum(bf16_efficient_factor) / len(bf16_efficient_factor)
+system_template["accelerator"]['op']['default'] = {
+ 'tflops' : MAX_TFLOPS,
+ 'efficient_factor': avg_efficient_factor,
+}
+json.dump(system_template, open(f"{sys_name}.json", "w"), indent=4)
\ No newline at end of file
diff --git a/simu_tools/efficency_test/nccl_fit.py b/simu_tools/efficency_test/nccl_fit.py
new file mode 100644
index 0000000..58c19cd
--- /dev/null
+++ b/simu_tools/efficency_test/nccl_fit.py
@@ -0,0 +1,132 @@
+from matplotlib import pyplot as plt
+import numpy as np
+from matplotlib import pyplot as plt
+import numpy as np
+
+DEVICE_BW = 900 # H100机内标称带宽
+'''
+https://www.advancedclustering.com/wp-content/uploads/2022/03/gtc22-whitepaper-hopper.pdf
+ While the 4-
+ GPU configuration includes point-to-point NVLink connections between GPUs and provides a
+ higher CPU-to-GPU ratio in the server, the 8-GPU configuration includes NVSwitch to provide
+ SHARP in-network reductions and full NVLink bandwidth of 900 GB/s between any pair of
+ GPUs. The H100 SXM5 GPU is also used in the powerful new DGX H100 servers and DGX
+ SuperPOD systems
+'''
+
+def get_bws(data, scale=1, ranks=8, title=""):
+ x = []
+ data = data.strip().split('\n') # 移除头尾空行
+ for d in data:
+ t = d.split()
+ n, cost, bw1, bw2 = t[0], t[5], t[6], t[7]
+ x.append([n, cost, bw1, bw2])
+ y = np.array(x, dtype=float)
+
+ # X: 通信量（B）
+ n = y[:, 0] * (ranks-1) / ranks * scale
+ cost = y[:, 1]
+
+ # 拟合
+ a, b = np.polyfit(n[:11], cost[:11], 1)
+ show = False
+ if show:
+ # -------- 第1张图：性能数据 + 拟合 --------
+ plt.figure(figsize=(8, 4))
+ plt.plot(n, cost, marker='o', label='Data')
+ plt.plot(n, n*a + b, '--', label='Linear Fit')
+ plt.xscale('log', base=2)
+ plt.xlabel("Bytes")
+ plt.ylabel("Time (us)")
+ plt.title(f"Cost vs Bytes- {title}")
+ plt.legend()
+ plt.grid(True, which='both', linestyle='--', alpha=0.5)
+
+ # -------- 第2张图：相对误差 --------
+ plt.figure(figsize=(8, 4))
+ rel_error = np.abs(n*a + b - cost) / cost
+ plt.plot(n, rel_error, marker='x', color='red')
+ plt.xscale('log', base=2)
+ plt.xlabel("Bytes")
+ plt.ylabel("Relative Error")
+ plt.title(f"Relative Fit Error- {title}")
+ plt.grid(True, which='both', linestyle='--', alpha=0.5)
+
+ # -------- 返回估算带宽 --------
+ bw = 1 / a / 1024**3 * 1000**2
+ eff = round(bw / DEVICE_BW , 4)
+
+ latency_per_p2p = b / ((ranks-1) * scale)
+ print(f"{title} Estimated Bandwidth: {bw:.2f} GB/s, Efficiency: {eff:.4f}, Latency per P2P: {latency_per_p2p:.2f} μs")
+ return bw, eff, latency_per_p2p # 带宽 (GB/s), 起始开销（μs）
+
+get_bws(data=''' 1048576 524288 bfloat16 sum -1 46.24 22.68 39.69 0 45.89 22.85 39.99 0
+ 2097152 1048576 bfloat16 sum -1 45.71 45.88 80.29 0 45.82 45.77 80.10 0
+ 4194304 2097152 bfloat16 sum -1 57.93 72.41 126.71 0 58.22 72.04 126.06 0
+ 8388608 4194304 bfloat16 sum -1 86.55 96.92 169.61 0 85.32 98.32 172.05 0
+ 16777216 8388608 bfloat16 sum -1 125.9 133.24 233.17 0 126.7 132.46 231.81 0
+ 33554432 16777216 bfloat16 sum -1 203.3 165.09 288.90 0 203.3 165.03 288.80 0
+ 67108864 33554432 bfloat16 sum -1 328.4 204.37 357.65 0 327.8 204.74 358.30 0
+ 134217728 67108864 bfloat16 sum -1 579.2 231.72 405.51 0 579.9 231.45 405.04 0
+ 268435456 134217728 bfloat16 sum -1 1090.5 246.16 430.79 0 1089.8 246.31 431.04 0
+ 536870912 268435456 bfloat16 sum -1 2106.2 254.90 446.07 0 2107.3 254.76 445.84 0
+ 1073741824 536870912 bfloat16 sum -1 4022.4 266.94 467.15 0 4017.6 267.26 467.70 0
+ 2147483648 1073741824 bfloat16 sum -1 7924.8 270.98 474.22 0 7921.2 271.10 474.43 0
+ 4294967296 2147483648 bfloat16 sum -1 15736 272.95 477.66 0 15755 272.62 477.08 0
+ 8589934592 4294967296 bfloat16 sum -1 31450 273.13 477.98 0 31436 273.26 478.20 0''', scale = 2, ranks = 8, title = 'all_reduce_nccl_test')
+get_bws(data=''' 1048576 65536 bfloat16 none -1 47.14 22.24 19.46 0 45.43 23.08 20.20 0
+ 2097152 131072 bfloat16 none -1 45.69 45.90 40.16 0 45.76 45.83 40.10 0
+ 4194304 262144 bfloat16 none -1 47.23 88.80 77.70 0 50.69 82.74 72.40 0
+ 8388608 524288 bfloat16 none -1 48.87 171.67 150.21 0 49.08 170.90 149.54 0
+ 16777216 1048576 bfloat16 none -1 73.03 229.73 201.02 0 70.13 239.21 209.31 0
+ 33554432 2097152 bfloat16 none -1 128.2 261.71 228.99 0 125.2 267.99 234.50 0
+ 67108864 4194304 bfloat16 none -1 211.3 317.58 277.88 0 208.6 321.64 281.43 0
+ 134217728 8388608 bfloat16 none -1 386.4 347.34 303.92 0 384.5 349.08 305.45 0
+ 268435456 16777216 bfloat16 none -1 730.2 367.64 321.68 0 728.5 368.49 322.43 0
+ 536870912 33554432 bfloat16 none -1 1394.3 385.04 336.91 0 1387.1 387.04 338.66 0
+ 1073741824 67108864 bfloat16 none -1 2734.1 392.72 343.63 0 2707.8 396.54 346.97 0
+ 2147483648 134217728 bfloat16 none -1 5379.0 399.23 349.33 0 5323.3 403.42 352.99 0
+ 4294967296 268435456 bfloat16 none -1 10615 404.62 354.05 0 10460 410.60 359.28 0
+ 8589934592 536870912 bfloat16 none -1 21025 408.57 357.49 0 20741 414.15 362.38 0''', scale = 1, ranks = 8, title = 'all_gather_nccl_test')
+get_bws(data='''1048576 65536 bfloat16 sum -1 43.62 24.04 21.03 0 43.24 24.25 21.22 0
+ 2097152 131072 bfloat16 sum -1 44.87 46.74 40.90 0 44.02 47.64 41.69 0
+ 4194304 262144 bfloat16 sum -1 59.15 70.91 62.05 0 46.64 89.92 78.68 0
+ 8388608 524288 bfloat16 sum -1 54.41 154.17 134.90 0 53.33 157.30 137.64 0
+ 16777216 1048576 bfloat16 sum -1 75.44 222.41 194.60 0 192.4 87.19 76.29 0
+ 33554432 2097152 bfloat16 sum -1 117.8 284.86 249.25 0 116.5 288.03 252.02 0
+ 67108864 4194304 bfloat16 sum -1 205.3 326.96 286.09 0 202.1 332.07 290.56 0
+ 134217728 8388608 bfloat16 sum -1 385.5 348.20 304.68 0 380.8 352.51 308.44 0
+ 268435456 16777216 bfloat16 sum -1 726.0 369.73 323.52 0 724.0 370.75 324.41 0
+ 536870912 33554432 bfloat16 sum -1 1410.0 380.76 333.16 0 1404.3 382.29 334.51 0
+ 1073741824 67108864 bfloat16 sum -1 2758.3 389.27 340.61 0 2750.0 390.46 341.65 0
+ 2147483648 134217728 bfloat16 sum -1 5419.6 396.25 346.72 0 5403.0 397.46 347.78 0
+ 4294967296 268435456 bfloat16 sum -1 10584 405.82 355.09 0 10599 405.22 354.56 0
+ 8589934592 536870912 bfloat16 sum -1 20940 410.21 358.94 0 20909 410.83 359.47 0''', scale = 1, ranks = 8, title = 'reduce_scatter_nccl_test')
+get_bws(data=''' 1048576 65536 bfloat16 none -1 75.76 13.84 12.11 0 74.98 13.99 12.24 N/A
+ 2097152 131072 bfloat16 none -1 75.06 27.94 24.45 0 74.05 28.32 24.78 N/A
+ 4194304 262144 bfloat16 none -1 76.64 54.73 47.89 0 76.01 55.18 48.28 N/A
+ 8388608 524288 bfloat16 none -1 79.10 106.05 92.80 0 79.59 105.40 92.23 N/A
+ 16777216 1048576 bfloat16 none -1 82.85 202.51 177.19 0 82.91 202.36 177.07 N/A
+ 33554432 2097152 bfloat16 none -1 129.9 258.33 226.04 0 128.7 260.75 228.15 N/A
+ 67108864 4194304 bfloat16 none -1 220.1 304.96 266.84 0 214.5 312.84 273.74 N/A
+ 134217728 8388608 bfloat16 none -1 398.7 336.60 294.52 0 398.5 336.84 294.74 N/A
+ 268435456 16777216 bfloat16 none -1 761.2 352.66 308.58 0 754.9 355.61 311.16 N/A
+ 536870912 33554432 bfloat16 none -1 1461.0 367.46 321.52 0 1461.0 367.46 321.53 N/A
+ 1073741824 67108864 bfloat16 none -1 2822.3 380.44 332.89 0 2821.4 380.57 333.00 N/A
+ 2147483648 134217728 bfloat16 none -1 5520.5 389.00 340.38 0 5515.5 389.36 340.69 N/A
+ 4294967296 268435456 bfloat16 none -1 10937 392.69 343.60 0 10998 390.52 341.71 N/A
+ 8589934592 536870912 bfloat16 none -1 21730 395.31 345.89 0 21834 393.43 344.25 N/A''', scale = 1, ranks = 8, title = 'alltoall_nccl_test')
+get_bws(data=''' 1048576 524288 bfloat16 sum -1 15.49 67.71 67.71 0 1.62 646.87 646.87 N/A
+ 2097152 1048576 bfloat16 sum -1 15.01 139.70 139.70 0 1.51 1387.01 1387.01 N/A
+ 4194304 2097152 bfloat16 sum -1 15.69 267.32 267.32 0 1.48 2832.08 2832.08 N/A
+ 8388608 4194304 bfloat16 sum -1 16.79 499.62 499.62 0 1.43 5886.74 5886.74 N/A
+ 16777216 8388608 bfloat16 sum -1 23.09 726.70 726.70 0 1.44 11691.44 11691.44 N/A
+ 33554432 16777216 bfloat16 sum -1 37.97 883.69 883.69 0 1.43 23448.24 23448.24 N/A
+ 67108864 33554432 bfloat16 sum -1 61.50 1091.22 1091.22 0 1.70 39522.30 39522.30 N/A
+ 134217728 67108864 bfloat16 sum -1 110.2 1217.81 1217.81 0 1.44 93466.38 93466.38 N/A
+ 268435456 134217728 bfloat16 sum -1 203.9 1316.39 1316.39 0 1.39 193816.21 193816.21 N/A
+ 536870912 268435456 bfloat16 sum -1 389.7 1377.78 1377.78 0 1.43 374386.97 374386.97 N/A
+ 1073741824 536870912 bfloat16 sum -1 763.8 1405.86 1405.86 0 1.40 769156.03 769156.03 N/A
+ 2147483648 1073741824 bfloat16 sum -1 1508.7 1423.37 1423.37 0 1.46 1472896.88 1472896.88 N/A
+ 4294967296 2147483648 bfloat16 sum -1 2999.7 1431.81 1431.81 0 1.63 2639807.80 2639807.80 N/A
+ 8589934592 4294967296 bfloat16 sum -1 5968.9 1439.11 1439.11 0 2.50 3435973.84 3435973.84 N/A''', scale = 1, ranks = 2, title = 'sendrecv_nccl_test')
\ No newline at end of file
diff --git a/simu_tools/efficency_test/nccl_test.sh b/simu_tools/efficency_test/nccl_test.sh
new file mode 100644
index 0000000..8f68e82
--- /dev/null
+++ b/simu_tools/efficency_test/nccl_test.sh
@@ -0,0 +1,11 @@
+end=8G
+echo "run all_reduce_perf"
+./build/all_reduce_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_all_reduce.txt
+echo "run all_gather_perf"
+./build/all_gather_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_all_gather.txt
+echo "run reduce_scatter_perf"
+./build/reduce_scatter_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_reduce_scatter.txt
+echo "run alltoall_perf"
+./build/alltoall_perf -n 10 -b 1M -e $end -f 2 -g 8 -w 2 -d bfloat16 > perf_alltoall_perf.txt
+echo "run sendrecv_perf"
+./build/sendrecv_perf -n 1 -b 1M -e $end -f 2 -g 1 -t 1 -d bfloat16 > send_recv.txt
\ No newline at end of file
diff --git a/simu_tools/efficency_test/reduce_scatter.py b/simu_tools/efficency_test/reduce_scatter.py
new file mode 100644
index 0000000..7718d45
--- /dev/null
+++ b/simu_tools/efficency_test/reduce_scatter.py
@@ -0,0 +1,51 @@
+import os
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.profiler import profile, record_function, ProfilerActivity
+
+def run_reduce_scatter(rank, world_size):
+ dist.init_process_group("nccl", rank=rank, world_size=world_size)
+ torch.cuda.set_device(rank)
+
+ # 每个rank准备 world_size 份输入，每份大小为 240MB / world_size
+ total_bytes = 240 * 1024 * 1024
+ chunk_bytes = total_bytes
+ chunk_elements = chunk_bytes // 2 # float32 每个 4 字节
+
+ input_list = [torch.randn(chunk_elements, dtype=torch.bfloat16).cuda(rank) for _ in range(world_size)]
+ output_tensor = torch.zeros(chunk_elements, dtype=torch.bfloat16).cuda(rank)
+
+ # Warmup
+ for _ in range(10):
+ dist.reduce_scatter(output_tensor, input_list, op=dist.ReduceOp.SUM)
+ torch.cuda.synchronize()
+
+ # 正式测试 + profiling
+ with profile(
+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+ schedule=torch.profiler.schedule(wait=0, warmup=5, active=50),
+ on_trace_ready=torch.profiler.tensorboard_trace_handler(f"./log_rank{rank}"),
+ record_shapes=True,
+ with_stack=True,
+ ) as prof:
+ for _ in range(50):
+ with record_function("reduce_scatter_comm"):
+ dist.reduce_scatter(output_tensor, input_list, op=dist.ReduceOp.SUM)
+ prof.step()
+ torch.cuda.synchronize()
+ dist.destroy_process_group()
+
+def main():
+ os.environ['MASTER_ADDR'] = '127.0.0.1'
+ os.environ['MASTER_PORT'] = '29500'
+ os.environ['RANK'] = '0'
+ os.environ['WORLD_SIZE'] = str(torch.cuda.device_count())
+
+ world_size = torch.cuda.device_count()
+ if world_size < 2:
+ raise RuntimeError("需要至少 2 张 GPU 才能进行 reduce-scatter 测试")
+ mp.spawn(run_reduce_scatter, args=(world_size,), nprocs=world_size, join=True)
+
+if __name__ == "__main__":
+ main()
diff --git a/simu_tools/efficency_test/run.sh b/simu_tools/efficency_test/run.sh
new file mode 100644
index 0000000..e96f140
--- /dev/null
+++ b/simu_tools/efficency_test/run.sh
@@ -0,0 +1,9 @@
+export MAX_TFLOPS=312 # 指定机器的标称算力
+export SYS_NAME="A100_PCIE" #指定system名称
+export PICE_INTRA_LINK=1 #指定机内卡间是否为PCIE互联
+export FC8_MODE=0 #指定机内卡间通信连接是否为FC8模式
+export PARAM_FILE="./run_params.json" # 测试超参数，指定测试模型的列表，mbs区间，seq_len区间等, 如果不指定，则使用默认超参。
+python test_gemm_efficency.py
+python test_grouped_gemm_efficency.py
+python test_fa_efficency.py
+python combine_efficency.py
\ No newline at end of file
diff --git a/simu_tools/efficency_test/run_params.json b/simu_tools/efficency_test/run_params.json
new file mode 100644
index 0000000..57eea54
--- /dev/null
+++ b/simu_tools/efficency_test/run_params.json
@@ -0,0 +1,16 @@
+{
+ "model_list":["deepseekv2",
+ "deepseekv3",
+ "deepseek-32b",
+ "deepseek-16b",
+ "deepseek-1b",
+ "llama3-8b",
+ "llama3-70b",
+ "qwen3-32b",
+ "kimi-1T",
+ "mixtral-8x7b"],
+ "mbs_list": [1, 2, 4],
+ "seq_len_list":[4096],
+ "tp_list": [1, 2, 4, 8],
+ "ep_list": [1, 2, 4, 8, 16, 64]
+}
\ No newline at end of file
diff --git a/simu_tools/efficency_test/test_fa_efficency.py b/simu_tools/efficency_test/test_fa_efficency.py
new file mode 100644
index 0000000..cbf45a0
--- /dev/null
+++ b/simu_tools/efficency_test/test_fa_efficency.py
@@ -0,0 +1,219 @@
+import time,os
+import torch
+from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward
+from simu_tools.efficency_test.utils import get_system_name, sync_device, get_torch_profiler, get_all_test_model_configs, get_test_seq_len_list, get_test_mbs_list, get_test_tp_list
+
+# 1. 构造测试输入数据
+def generate_test_inputs(device, batch_size, seq_len, num_q_heads, num_kv_heads, k_head_dim, v_head_dim, qkv_contiguous):
+ # batch_size, seq_len, num_heads, head_dim = 4096, 1, 4, 128
+ # 生成随机输入张量 (使用BF16格式以测试FlashAttention)
+
+ query = torch.randn((batch_size, seq_len, num_q_heads, k_head_dim),
+ dtype=torch.bfloat16, device=device)
+ key = torch.randn((batch_size, seq_len, num_kv_heads, k_head_dim),
+ dtype=torch.bfloat16, device=device)
+ value = torch.randn((batch_size, seq_len, num_kv_heads, v_head_dim),
+ dtype=torch.bfloat16, device=device)
+
+ return query, key, value
+
+# 4. 基准测试函数
+def benchmark_flashattention(q, k, v, warmup=10, repeat=100, device='cuda'):
+ dq = torch.empty_like(q)
+ dk = torch.empty_like(k)
+ dv = torch.empty_like(v)
+ q.requires_grad_()
+ k.requires_grad_()
+ v.requires_grad_()
+ fwd_elapsed_time, bwd_elapsed_time = 0, 0
+ profiler = get_torch_profiler(device, False)
+ # Benchmark
+ sync_device(device)
+ for i in range(repeat):
+ if i == warmup:
+ sync_device(device)
+ fwd_start_time = time.time()
+ # with torch.no_grad():
+ with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
+ # fwd
+ if device == 'cuda':
+ out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+ q = q,
+ k = k,
+ v = v,
+ dropout_p = 0.0,
+ softmax_scale = 0.1,
+ causal = True,
+ window_size = (-1,0),
+ alibi_slopes = None,
+ return_softmax = False
+ )
+ else:
+ raise ValueError(f"Unknown device: {device}")
+ profiler.step() if profiler else 1
+
+ sync_device(device)
+ fwd_elapsed_time = (time.time() - fwd_start_time) / (repeat-warmup) * 1000
+
+ # print(f"=== FlashAttention输出: {attn_output}")
+ # for i, o in enumerate(attn_output):
+ # print(f"=== 输出{i}: {type(o)}")
+ # 测试反向传播
+
+ sync_device(device)
+ for i in range(repeat):
+ if i == warmup:
+ sync_device(device)
+ bwd_start_time = time.time()
+ # with torch.no_grad():
+ if 1:
+ with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
+ if device == 'cuda':
+ grad_input = _flash_attn_backward(
+ dout = out,
+ # dump_grad,
+ q = q,
+ k = k,
+ v = v, #q, k, v,
+ out = out,
+ softmax_lse = softmax_lse,
+ dq = dq,
+ dk = dk,
+ dv = dv,
+ dropout_p=0.0,
+ softmax_scale = 1.0,
+ causal = True,
+ window_size=(-1,0),
+ alibi_slopes=None,
+ deterministic=False
+ # *detached_ori_outputs, #(Tensor output, Tensor logsumexp, Tensor dropout_mask)
+ # is_causal=True #causal same as fwd
+ )
+ else:
+ raise ValueError(f"Unknown device: {device}")
+ sync_device(device)
+ bwd_elapsed_time = (time.time() - bwd_start_time) / (repeat - warmup) * 1000
+
+ profiler.stop() if profiler else 1
+ return fwd_elapsed_time, bwd_elapsed_time
+
+def test(model, all_efficiency, qkv_contiguous, batch, seq_len, num_q_heads, num_kv_heads, qk_head_dim, v_head_dim, TP=1, MAX_TFLOPS=None, res=None, device='cuda'):
+ assert num_q_heads % TP == 0 and num_kv_heads % TP == 0
+ num_q_heads //= TP
+ num_kv_heads //= TP
+ shape_key = f'batch={batch}, seq_len={seq_len}, head_num={num_q_heads}, kv_head_num={num_kv_heads}, qk_head_dim={qk_head_dim}, v_head_dim={v_head_dim}, qkv_contiguous={qkv_contiguous}'
+ if shape_key in all_efficiency['sdp_fwd']['accurate_efficient_factor']:
+ return
+ query, key, value = generate_test_inputs(device, batch, seq_len, num_q_heads, num_kv_heads, qk_head_dim, v_head_dim, qkv_contiguous)
+ print(f"- 输入形状: query={query.shape}, key={key.shape}, value={value.shape}")
+ print(f'batch: {batch}, seq_len: {seq_len}, num_heads: {num_q_heads}&{num_kv_heads}, head_dim: {qk_head_dim}')
+ # 初始化配置和模块
+ # 运行基准测试
+ fwd_latency, bwd_latency = benchmark_flashattention(query, key, value, device = device)
+
+ # 计算理论FLOPs (公式: 2 * batch * seq_len^2 * num_heads * head_dim)
+ base_flops = batch * (seq_len ** 2) * max(num_q_heads, num_kv_heads) * (qk_head_dim + v_head_dim)
+ fwd_flops = 2 * base_flops
+ fwd_tflops = (fwd_flops / (fwd_latency * 1e-3+ 1e-12)) / 1e12 # 转换为TFLOPs
+
+ bwd_flops = 5 * base_flops
+ bwd_tflops = (bwd_flops / (bwd_latency * 1e-3+ 1e-12)) / 1e12 # 转换为TFLOPs
+
+ print(f"=== {model} FlashAttention性能结果(TP={TP}):")
+ print(f"- 延迟: {fwd_latency:.3f} ms/iteration, backward: {bwd_latency:.3f} ms/iteration")
+ print(f"- fwd吞吐量: {fwd_tflops:.2f} TFLOPs, flops={fwd_flops} latency={fwd_latency:.2f} ms, 计算效率={fwd_tflops/MAX_TFLOPS:.2f}")
+ print(f"- bwd吞吐量: {bwd_tflops:.2f} TFLOPs, flops={bwd_flops} latency={bwd_latency:.2f} ms, 计算效率={bwd_tflops/MAX_TFLOPS:.2f}")
+ print(f"- 输入形状: query={query.shape}, key={key.shape}, value={value.shape}")
+ res['model'].append(model)
+ res['TP'].append(TP)
+ res['batch'].append(batch)
+ res['seq_len'].append(seq_len)
+ res['num_heads'].append(f'q={num_q_heads}, kv={num_kv_heads}')
+ res['head_size'].append(f'qk={qk_head_dim//TP:.0f}, v={v_head_dim//TP:.0f}')
+ res['flops'].append(f'fwd={fwd_flops}, bwd={bwd_flops}')
+ res['time'].append(f'fwd={fwd_latency:.2f} ms, bwd={bwd_latency:.2f} ms')
+ res['TFLOPS'].append(f'fwd={fwd_tflops:.2f} TFLOPS, bwd={bwd_tflops:.2f} TFLOPS')
+ res['fwd_efficiency'].append(round(fwd_tflops/MAX_TFLOPS, 2))
+ res['bwd_efficiency'].append(round(bwd_tflops/MAX_TFLOPS, 2))
+
+
+ sdp_key = 'sdp_fwd'
+ all_efficiency[sdp_key]['accurate_efficient_factor'][shape_key] = fwd_tflops/MAX_TFLOPS
+ all_efficiency[sdp_key]['efficient_factor'] = sum(all_efficiency[sdp_key]['accurate_efficient_factor'].values()) / len(all_efficiency[sdp_key]['accurate_efficient_factor'])
+
+ sdp_key = 'sdp_bwd'
+ all_efficiency[sdp_key]['accurate_efficient_factor'][shape_key] = bwd_tflops/MAX_TFLOPS
+ all_efficiency[sdp_key]['efficient_factor'] = sum(all_efficiency[sdp_key]['accurate_efficient_factor'].values()) / len(all_efficiency[sdp_key]['accurate_efficient_factor'])
+
+
+
+# 5. 主测试流程
+if __name__ == "__main__":
+ # 准备输入数据
+ system, device, MAX_TFLOPS = get_system_name()
+ save_root = f'{system}_fa_efficency'
+ os.makedirs(save_root, exist_ok=True)
+ MODEL_CONFIGS = get_all_test_model_configs()
+ MBS_LIST = get_test_mbs_list()
+ SEQ_LEN_LIST = get_test_seq_len_list()
+ TP_LIST = get_test_tp_list()
+ res = {
+ 'model':[],
+ 'TP':[],
+ 'batch':[],
+ 'seq_len':[],
+ 'num_heads':[],
+ 'head_size':[],
+ 'flops':[],
+ 'TFLOPS':[],
+ 'time':[],
+ 'fwd_efficiency':[],
+ 'bwd_efficiency':[],
+ # 'shape_str':[]
+ }
+ merged_dict = {
+ 'sdp_fwd':{
+ 'tflops': MAX_TFLOPS,
+ 'efficient_factor': 0,
+ 'accurate_efficient_factor':{
+ }
+ },
+ 'sdp_bwd':{
+ 'tflops': MAX_TFLOPS,
+ 'efficient_factor': 0,
+ 'accurate_efficient_factor':{
+ }
+ }
+ }
+ for SEQ_LEN in SEQ_LEN_LIST:
+ for MBS in MBS_LIST:
+ for model_config in MODEL_CONFIGS:
+ for tp in TP_LIST:
+ model = model_config.model_name
+ print(f"Running {model}...")
+ if model_config.qk_head_dim is not None:
+ qk_head_dim = model_config.qk_head_dim + (model_config.qk_pos_emb_head_dim if model_config.qk_pos_emb_head_dim is not None else 0)
+ else:
+ qk_head_dim = model_config.head_size
+ v_head_dim = model_config.v_head_dim if model_config.v_head_dim is not None else model_config.head_size
+
+ params = dict(batch=MBS,
+ seq_len=SEQ_LEN,
+ num_q_heads = model_config.head_num,
+ num_kv_heads = model_config.kv_head_num,
+ qk_head_dim = qk_head_dim,
+ v_head_dim = v_head_dim,
+ TP = tp,
+ MAX_TFLOPS = MAX_TFLOPS,
+ res = res,
+ device = device)
+ test(model, merged_dict, True, **params)
+ if device == 'cuda':
+ test(model, merged_dict, False, **params)
+
+ import pandas as pd
+ import json
+ df = pd.DataFrame(res)
+ df.to_csv(os.path.join(save_root, 'fa_efficency_test.csv'), index=False)
+ with open(os.path.join(save_root, 'fa_efficency_test.json'), 'w') as f:
+ json.dump(merged_dict, f, indent=4)
diff --git a/simu_tools/efficency_test/test_gemm_efficency.py b/simu_tools/efficency_test/test_gemm_efficency.py
new file mode 100644
index 0000000..b274483
--- /dev/null
+++ b/simu_tools/efficency_test/test_gemm_efficency.py
@@ -0,0 +1,293 @@
+import os
+import time
+import json
+from argparse import ArgumentParser
+import pandas as pd
+from matplotlib import pyplot as plt
+import torch
+from simumax.core.config import ModelConfig, StrategyConfig, SystemConfig
+from simumax.core.perf_llm import PerfLLM
+import transformer_engine as te
+import transformer_engine_torch as tex
+from transformer_engine.pytorch.module.base import get_workspace
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
+from transformer_engine.pytorch.cpp_extensions import (
+ general_gemm
+)
+from simumax.utils import get_simu_model_config, get_simu_system_config, get_simu_strategy_config
+from simu_tools.efficency_test.utils import get_system_name, sync_device, get_torch_profiler, get_all_test_model_configs, get_test_seq_len_list, get_test_mbs_list, get_test_tp_list, get_test_ep_list
+
+def prepare_tensors(B: int, M: int, K: int, N: int, dtype, device):
+ """
+ 准备输入矩阵和workspace
+ Args:
+ B: batch size
+ M: A矩阵行数
+ N: B矩阵列数
+ K: A矩阵列数/B矩阵行数
+ Returns:
+ A: shape [B, M, K] (TN布局)
+ B: shape [K, N] (TN布局)
+ workspace: 预分配空间
+ """
+ if dtype == 'bf16':
+ dtype = torch.bfloat16
+ else:
+ raise ValueError("dtype must be 'bf16'")
+
+ A = torch.randn((B, M, K), device=device, dtype=dtype)
+ B_matrix = torch.randn((K, N), device=device, dtype=dtype)
+
+ # 为general_gemm分配workspace (典型大小32MB)
+ workspace = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device=device)
+
+ return A, B_matrix, workspace
+
+def run_te_gemm(B: int, M: int, K: int, N: int, dtype, device, layout, accumulate, out_dtype, test_steps, warmup_steps):
+ """
+ 执行general_gemm操作
+ Args:
+ A: [B, M, K]
+ B: [K, N]
+ Returns:
+ output: [B, M, N]
+ """
+
+ out_dtype = torch.bfloat16 if out_dtype == 'bf16' else torch.float32
+ out_dtype = torch.bfloat16
+
+ if layout == 'NN':
+ B_matrix = torch.randn((K, N), device=device, dtype=torch.bfloat16)
+ A_matrix = torch.randn((M, B, K), device=device, dtype=torch.bfloat16)
+ output = torch.empty((M, B, N), device=device, dtype=out_dtype)
+ elif layout == 'TN':
+ B_matrix = torch.randn((N, K), device=device, dtype=torch.bfloat16)
+ A_matrix = torch.randn((M, B, K), device=device, dtype=torch.bfloat16)
+ output = torch.empty((M, B, N), device=device, dtype=out_dtype)
+ elif layout == 'NT':
+ B_matrix = torch.randn((K, B, N), device=device, dtype=torch.bfloat16) # N
+ A_matrix = torch.randn((K, B, M), device=device, dtype=torch.bfloat16) # T
+ output = torch.empty((M, N), device=device, dtype=out_dtype)
+ else:
+ raise ValueError("layout must be 'NN' or 'TN' or 'NN'")
+ # 为general_gemm分配workspace (典型大小32MB)
+ print(f'dtype: {dtype}, out_dtype: {out_dtype}, layout: {layout}')
+ if dtype == 'fp8':
+ # quantizer = Float8Quantizer(torch.tensor(1.0), amax=torch.tensor(1.0), fp8_dtype=tex.DType.kBFloat16)
+ te_dtype = tex.DType.kFloat8E4M3
+ quantizer = Float8Quantizer(
+ scale=torch.full([1], 1.0).to(device).squeeze(),
+ amax=torch.full([1], 1.0).to(device),
+ fp8_dtype=te_dtype,
+ )
+ quantizer.set_usage(rowwise=True, columnwise=False)
+ A_matrix = quantizer(A_matrix)
+ B_matrix = quantizer(B_matrix)
+ print(f'A_matrix: {A_matrix.shape} {A_matrix.device} {A_matrix.dtype}, B_matrix: {B_matrix.shape} {B_matrix.device} {B_matrix.dtype}, output: {output.shape} {output.device} {output.dtype}')
+
+ assert warmup_steps < test_steps, "warmup_steps should be less than test_steps"
+ print(f'B={B_matrix.shape}, A={A_matrix.shape}, out={output.shape}, layout={layout}')
+ sync_device(device)
+
+ for i in range(test_steps):
+ # 处理batch维度：对每个batch元素执行GEMM
+ if i == warmup_steps:
+ sync_device(device)
+ start_time = time.time()
+ general_gemm(
+ B_matrix,
+ A_matrix,
+ get_workspace(),
+ out_dtype=out_dtype,
+ quantization_params=None,
+ out=output,
+ accumulate=accumulate,
+ layout=layout,
+ grad = True if layout in ['NT', 'NN'] else False,
+ use_split_accumulator = True if layout in ['NT', 'NN'] else False,
+ )
+ """
+ wgrad, grad_bias_, _, rs_out = ceg.general_gemm(
+ inputmat_total,
+ grad_output,
+ get_workspace(),
+ layout="NT",
+ grad=True,
+ out_dtype=(
+ main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype
+ ),
+ bias=(bias if (grad_bias is None and not ctx.fp8) else None),
+ out=main_grad if ctx.fuse_wgrad_accumulation else None,
+ use_split_accumulator=_2X_ACC_WGRAD,
+ accumulate=accumulate_wgrad_into_param_main_grad,
+ ub=ub_obj_wgrad,
+ ub_type=ub_type_wgrad,
+ extra_output=rs_out,
+ bulk_overlap=ctx.ub_bulk_wgrad,
+ )
+
+ """
+ # 检查output中是否存在-1e10，如果存在则说明有错误
+ # if torch.any(output == -1e10):
+ # raise ValueError("general_gemm returned an invalid value")
+ sync_device(device)
+ end_time = time.time()
+ dur = (end_time - start_time) / (test_steps - warmup_steps)
+ return dur
+
+def plot_topk(ops_info, topk, save_path):
+ plt.cla()
+ plt.pie(ops_info['percentage'].head(topk), labels=ops_info['Module'].head(topk), autopct='%1.1f%%')
+ plt.savefig(save_path)
+ plt.show()
+
+def test_gemm_efficency(gemm_shape_list, max_tflops, device, save_root, dtype, res, test_steps, warmup_steps):
+ # 遍历df每一行，拿到b, m n k， 测试gemm shape的效率，记录在efficency中
+ efficiency_file_path = f'{save_root}/gemm_efficency.json'
+ matmul_key = 'matmul' if dtype == 'bf16' else 'fp8_matmul'
+ if os.path.exists(efficiency_file_path):
+ all_efficiency = json.load(open(efficiency_file_path))
+ else:
+ all_efficiency = {
+ 'matmul': {
+ 'tflops': max_tflops,
+ 'efficient_factor': 0,
+ 'accurate_efficient_factor':{
+ }
+ },
+ 'fp8_matmul': {
+ 'tflops': max_tflops,
+ 'efficient_factor': 0,
+ 'accurate_efficient_factor':{
+ }
+ }
+ }
+ accurate_efficiency:dict = all_efficiency[matmul_key]['accurate_efficient_factor']
+
+ # Assuming df is a pandas DataFrame with columns 'b', 'm', 'n', 'k'
+ for index, row in gemm_shape_list.iterrows():
+ b = row['B'] # batch size
+ m = row['M'] # rows of matrix A
+ k = row['K'] # columns of matrix A / rows of matrix B
+ n = row['N'] # columns of matrix B
+ layout = row['layout']
+ accumulate = row['accumulate']
+ out_dtype = row['out_dtype']
+
+ # Key for the efficiency dictionary - you can use a tuple as the key
+ shape_key = f'b={b}, m={m}, k={k}, n={n}, layout={layout}, accumulate={accumulate}, out_dtype={out_dtype}'
+
+ # Here you would perform your GEMM operation and measure its efficiency
+ # This is a placeholder - replace with actual measurement code
+ # For example, you might time the operation or measure FLOPs
+ if accurate_efficiency.get(shape_key, None) is None:
+ # Pseudocode for measurement:
+ # Perform GEMM operation with shape (b, m, k, n)
+ # A, B_matrix, workspace = prepare_tensors(b, m, k, n, dtype, device)
+ execution_time = run_te_gemm(b, m, k, n,dtype,
+ device=device,
+ layout = layout,
+ accumulate=accumulate,
+ out_dtype = out_dtype,
+ test_steps=test_steps,
+ warmup_steps=warmup_steps)
+
+ # Calculate efficiency (this depends on your metric - could be GFLOPS, time, etc.)
+ # For example, theoretical FLOPs for GEMM: 2 * b * m * n * k
+ flops = 2 * b * m * n * k
+ tflops = flops / (execution_time * 1e12) if execution_time > 0 else 0
+
+ # Store the efficiency metric
+ accurate_efficiency[shape_key] = tflops/max_tflops
+
+ print(f"Shape {shape_key}: {tflops:.2f} TFLOPS, efficiency: {tflops/max_tflops}, execution_time={execution_time*1000} ms")
+ res[shape_key] = {'efficiency': tflops/max_tflops, 'execution_time' : f'{execution_time*1000} ms'}
+ else:
+ print(f'Shape {shape_key} already exists')
+
+ avg_efficiency = sum(accurate_efficiency.values())/len(accurate_efficiency)
+ all_efficiency[matmul_key]['efficient_factor'] = avg_efficiency
+ print(f'avg_efficiency={avg_efficiency}')
+
+ with open(efficiency_file_path, 'w') as f:
+ json.dump(all_efficiency, f, indent=4)
+
+def parse_ops_info(perf_model:PerfLLM, save_root):
+ # model_config = perf_model.model_config
+ # strategy_config = perf_model.strategy
+ op_infos = perf_model.analysis_op_info()
+ from pprint import pprint
+ df = pd.DataFrame(op_infos['first_stage_chunk'])
+ df['percentage'] = df['cost']/df['cost'].sum()
+ # stage = 'first_stage_chunk'
+ # df.to_csv(f"{save_root}/{model_config.model_type}_mbs_{strategy_config.micro_batch_size}_tp_{strategy_config.tp_size}_op_info_{stage}.csv", index=False)
+ # plot_topk(df, 20, f"{save_root}/{model_config.model_type}_mbs_{strategy_config.micro_batch_size}_tp_{strategy_config.tp_size}_op_info_{stage}_topk_op.png")
+ # print(df.head(100))
+ return df
+
+def test(dtype, grad_reduce_in_bf16):
+ system, device, MAX_TFLOPS = get_system_name()
+ save_root = f'{system}_gemm_efficency'
+ os.makedirs(save_root, exist_ok=True)
+ MODEL_CONFIGS = get_all_test_model_configs()
+ MBS_LIST = get_test_mbs_list()
+ SEQ_LEN_LIST = get_test_seq_len_list()
+
+ res = {}
+ system_config = SystemConfig.init_from_config_file(get_simu_system_config('a100_pcie'))
+ strategy_config = StrategyConfig.init_from_format_strings("tp1.ep1.tp1.pp1.mbs1.gbs8.world_size8")
+ strategy_config.fp8 = False
+ strategy_config.grad_reduce_in_bf16 = grad_reduce_in_bf16
+ perf_model = PerfLLM()
+
+ for SEQ_LEN in SEQ_LEN_LIST:
+ for model_config in MODEL_CONFIGS:
+ tp_list = [1]
+ ep_list = [1]
+ model_name = model_config.model_name
+ if model_config.model_type == 'moe':
+ model_config.moe_pad_expert_input_to_capacity = True
+ model_config.capacity = 1
+ model_config.layer_num = 10
+ elif model_config.model_type == 'dense':
+ tp_list = get_test_tp_list()
+
+ model_config.padded_vocab_size = True
+ model_config.make_vocab_size_divisible_by = 128
+
+ for MBS in MBS_LIST:
+ for tp in tp_list:
+ for ep in ep_list:
+ strategy_config.tp_size = tp
+ strategy_config.ep_size = ep
+ strategy_config.pp_size = 1
+ strategy_config.micro_batch_size = MBS # SET micro_batch_size
+ strategy_config.world_size = 8
+ strategy_config.seq_len = SEQ_LEN # SET sequence-length
+
+ perf_model.configure(
+ strategy_config=strategy_config,
+ model_config=model_config,
+ system_config=system_config
+ )
+ perf_model.run_estimate()
+ perf_model.analysis()
+
+ ops_info = parse_ops_info(perf_model, os.path.join(save_root, model_name))
+ test_gemm_efficency(gemm_shape_list = ops_info,
+ max_tflops = MAX_TFLOPS,
+ device = device,
+ save_root = save_root,
+ dtype = dtype,
+ res=res,
+ test_steps = 100,
+ warmup_steps = 25)
+
+ with open(os.path.join(save_root, 'all_efficiency_and_duration.json'), 'w') as f:
+ json.dump(res, f, indent=4)
+
+if __name__ == '__main__':
+ # test('fp8', grad_reduce_in_bf16=False)
+ # test('fp8', grad_reduce_in_bf16=True)
+ test('bf16', grad_reduce_in_bf16=False)
+ test('bf16', grad_reduce_in_bf16=True)
\ No newline at end of file
diff --git a/simu_tools/efficency_test/test_grouped_gemm_efficency.py b/simu_tools/efficency_test/test_grouped_gemm_efficency.py
new file mode 100644
index 0000000..0609f4a
--- /dev/null
+++ b/simu_tools/efficency_test/test_grouped_gemm_efficency.py
@@ -0,0 +1,269 @@
+import time
+import json
+import os
+import math
+from collections import OrderedDict
+import torch
+import transformer_engine as te
+import transformer_engine_torch as tex
+from transformer_engine.pytorch.module.base import get_multi_stream_cublas_workspace, _2X_ACC_FPROP
+from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
+from simu_tools.efficency_test.utils import get_system_name, sync_device, get_torch_profiler, get_all_test_model_configs, get_test_seq_len_list, get_test_mbs_list, get_test_ep_list
+
+torch_dtype_mapping = {
+ 'fp32': torch.float32,
+ 'bf16': torch.bfloat16,
+ 'torch.bfloat16' : 'bf16',
+ 'torch.float32' : 'fp32'
+}
+
+def run_te_grouped_gemm(A, B, out, m_splits, layout, grad,accumulate, use_split_accumulator, single_output, device, biases = None, warmup=10, repeat=50):
+ profiler = get_torch_profiler(device)
+ sync_device(device)
+ for i in range(repeat):
+ if i == warmup:
+ sync_device(device)
+ start_time = time.time()
+ _ = general_grouped_gemm(
+ A,
+ B,
+ out,
+ layout = layout,
+ out_dtype = out[0].dtype,
+ workspaces=get_multi_stream_cublas_workspace(),
+ m_splits=m_splits,
+ D_dtype = None,
+ use_bias=False,
+ bias = biases,
+ gelu = False,
+ grad = grad,
+ accumulate = accumulate,
+ use_split_accumulator = use_split_accumulator,
+ single_output = single_output
+ )
+ profiler.step() if profiler else 1
+ sync_device(device)
+ end_time = time.time()
+ return (end_time - start_time)/(repeat-warmup)*1000
+
+
+def run_grouped_gemm_fwd_bwd(ng, M, N, K, device, dtype, MAX_TFLOPS, grad_reduce_in_bf16, save_root):
+ # single_out_shape = ( # (M, N) x ng-> (ng * M, N) concat all experts' output
+ json_path = os.path.join(save_root, 'grouped_gemm_efficency.json')
+ all_efficiency = json.load(open(json_path, "r")) if os.path.exists(json_path) else {
+ 'group_matmul':{
+ 'tflops': MAX_TFLOPS,
+ 'efficient_factor': 0,
+ 'accurate_efficient_factor':{
+ }
+ },
+ 'fp8_group_matmul':{
+ 'tflops': MAX_TFLOPS,
+ 'efficient_factor': 0,
+ 'accurate_efficient_factor':{
+ }
+ }
+ }
+ matmul_key = 'fp8_group_matmul' if dtype == 'fp8' else 'group_matmul'
+ # multi_x = None
+ # multi_w = None
+ # multi_w_main_grad = None
+ # grad_output = None
+ # single_out = None
+ # m_splits = None
+
+ fwd_out_dtype = 'bf16'
+ main_grad_dtype = 'bf16' if grad_reduce_in_bf16 else 'fp32'
+ fwd_shape_str = f'ng={ng}, M={M}, N={N}, K={K}, dtype={dtype}, out_dtype={fwd_out_dtype}, main_grad_dtype={main_grad_dtype}, stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True'
+ bwd_grad_act_shape_str = f'ng={ng}, M={M}, N={N}, K={K}, dtype={dtype}, out_dtype={fwd_out_dtype}, main_grad_dtype={main_grad_dtype}, stage=bwd_grad_act, grad=True, accumulate=False, use_split_accumulator=True, single_output=False'
+ bwd_grad_weight_shape_str = f'ng={ng}, M={M}, N={N}, K={K}, dtype={dtype}, out_dtype={fwd_out_dtype}, main_grad_dtype={main_grad_dtype}, stage=bwd_grad_w, grad=True, accumulate=True, use_split_accumulator=True, single_output=False'
+
+ x_quantizer = Float8Quantizer(
+ scale=torch.full([1], 0.1).to(device).squeeze(),
+ amax=torch.full([1], 1.0).to(device),
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ )
+ w_quantizer = Float8Quantizer(
+ scale=torch.full([1], 0.12).to(device).squeeze(),
+ amax=torch.full([1], 1.0).to(device),
+ fp8_dtype=tex.DType.kFloat8E5M2,
+ )
+ def create_tensor():
+ multi_x = [torch.randn(M, K, dtype=torch.bfloat16, device=device) for _ in range(ng)]
+ multi_w = [torch.randn(N, K, dtype=torch.bfloat16, device=device) for _ in range(ng)]
+
+ single_out = [torch.randn(int(ng*M), N, dtype=torch.bfloat16, device=device)]
+
+ multi_w_main_grad = [torch.randn(N, K, dtype = (torch.float32 if main_grad_dtype == 'fp32' else torch.bfloat16), device=device) for _ in range(ng)]
+ grad_input = [torch.randn(M, N, dtype=torch.bfloat16, device=device) for _ in range(ng)]
+ grad_output = [torch.randn(M, K, dtype=torch.bfloat16, device=device) for _ in range(ng)]
+
+ m_splits = [M] * ng
+
+ if dtype == 'fp8':
+ multi_x = [x_quantizer(x) for x in multi_x]
+ multi_w = [w_quantizer(w) for w in multi_w]
+
+ return multi_x, multi_w, multi_w_main_grad, grad_input, grad_output, single_out, m_splits
+ if any([all_efficiency[matmul_key]['accurate_efficient_factor'].get(fwd_shape_str, None) is None,
+ all_efficiency[matmul_key]['accurate_efficient_factor'].get(bwd_grad_act_shape_str, None) is None,
+ all_efficiency[matmul_key]['accurate_efficient_factor'].get(bwd_grad_weight_shape_str, None) is None]):
+ multi_x, multi_w, multi_w_main_grad, grad_input, grad_output, single_out, m_splits = create_tensor()
+ else:
+ print(f'skip all!!!!!!!!!!!!')
+ return
+ # fwd
+ # groupgemm layout=TN, A=torch.Size([3072, 5120]) x 20, B=torch.Size([1232, 5120]) x 20, out=torch.Size([24640, 3072]) x 1, single_output=True, accumulate=False, grad=False, gelu=False, out_dtype=torch.bfloat16, D_dtype=None, use_bias=False, use_split_accumulator=False
+
+
+ if all_efficiency[matmul_key]['accurate_efficient_factor'].get(fwd_shape_str, None) is None:
+ fwd_dur = run_te_grouped_gemm(
+ A=multi_w[::-1], # ng x [N, K]
+ B=multi_x, # ng x [M, K] -> ng x [M, K] * [K, N] = ng x [M, N]
+ out =single_out,
+ m_splits=m_splits,
+ layout='TN',
+ grad=False,
+ accumulate=False,
+ use_split_accumulator=False,
+ single_output=True,
+ device=device)
+ flops = 2 * M * N * K * ng
+ fwd_TFLOPS = flops / (fwd_dur* 1e-3 + 1e-12) / 1e12
+ fwd_efficency = fwd_TFLOPS / MAX_TFLOPS
+
+ all_efficiency[matmul_key]['accurate_efficient_factor'][fwd_shape_str] = fwd_efficency
+ print(f"-- [fwd] Running grouped gemm with ng={ng}, M={M}, N={N}, K={K}, dtype={dtype}, out_dtype={fwd_out_dtype}, main_grad_dtype={main_grad_dtype}, device={device}, TFLOPS={fwd_TFLOPS}, dur={fwd_dur} ms, fwd_efficency={fwd_efficency}")
+ else:
+ print(f'-- [fwd] skip {fwd_shape_str}')
+ # bwd_grad_act
+ # groupgemm layout=NN, A=torch.Size([3072, 5120]) x 20, B=torch.Size([1232, 3072]) x 20, out=torch.Size([1232, 5120]) x 20, single_output=False, accumulate=False, grad=True, gelu=False, out_dtype=torch.bfloat16, D_dtype=None, use_bias=False, use_split_accumulator=True
+ # bwd_grad_act
+ if all_efficiency[matmul_key]['accurate_efficient_factor'].get(bwd_grad_act_shape_str, None) is None:
+ if dtype == 'fp8':
+ fp8_grad_input = [x_quantizer(g) for g in grad_input]
+ bwd_grad_act_dur = run_te_grouped_gemm(
+ A = multi_w, # ng x [N, K]
+ B = fp8_grad_input if dtype == 'fp8' else grad_input, # ng x [M, N] -> ng x [M, N]* [N , K] = ng x [M, K]
+ out = grad_output,
+ m_splits=m_splits,
+ layout='NN',
+ grad=True,
+ accumulate=False,
+ use_split_accumulator=True,
+ single_output=False,
+ device=device
+ )
+ flops = 2 * M * N * K * ng
+ bwd_grad_act_TFLOPS = flops / (bwd_grad_act_dur* 1e-3 + 1e-12) / 1e12
+ bwd_grad_act_efficency = bwd_grad_act_TFLOPS / MAX_TFLOPS
+ all_efficiency[matmul_key]['accurate_efficient_factor'][bwd_grad_act_shape_str] = bwd_grad_act_efficency
+ print(f"-- [bwd_grad_act] Running grouped gemm with ng={ng}, M={M}, N={N}, K={K}, dtype={dtype}, out_dtype={fwd_out_dtype}, main_grad_dtype={main_grad_dtype}, device={device}, TFLOPS={bwd_grad_act_TFLOPS}, dur={bwd_grad_act_dur} ms, bwd_grad_act_efficency={bwd_grad_act_efficency}")
+ else:
+ print(f'-- [bwd_grad_act] skip {bwd_grad_act_shape_str}')
+
+ # bwd_grad_weight
+ # groupgemm layout=NT, A=torch.Size([1232, 5120]) x 20, B=torch.Size([1232, 3072]) x 20, out=torch.Size([3072, 5120]) x 20, single_output=False, accumulate=True, grad=True, gelu=False, out_dtype=torch.bfloat16, D_dtype=None, use_bias=False, use_split_accumulator=True
+
+ if all_efficiency[matmul_key]['accurate_efficient_factor'].get(bwd_grad_weight_shape_str, None) is None:
+ if dtype == 'fp8':
+ fp8_grad_input2 = [x_quantizer(g) for g in grad_input]
+ biases = [torch.tensor([], dtype=torch.bfloat16, device=device) for _ in range(ng)]
+ else:
+ biases = None
+ bwd_grad_w_dur = run_te_grouped_gemm(
+ A= multi_x, # ng x [M, K]
+ B= fp8_grad_input2 if dtype == 'fp8' else grad_input, # ng x [M, N] -> ng x [N, M] * ng x [M, K] = ng x [N, K]
+ out = multi_w_main_grad,
+ m_splits=m_splits,
+ layout='NT',
+ grad=True,
+ accumulate=True,
+ use_split_accumulator=True,
+ single_output=False,
+ device=device,
+ biases = biases
+ )
+ flops = 2 * M * N * K * ng
+ bwd_grad_w_TFLOPS = flops / (bwd_grad_w_dur* 1e-3 + 1e-12) / 1e12
+ bwd_grad_w_efficency = bwd_grad_w_TFLOPS / MAX_TFLOPS
+ all_efficiency[matmul_key]['accurate_efficient_factor'][bwd_grad_weight_shape_str] = bwd_grad_w_efficency
+ print(f"-- [bwd_grad_w] Running grouped gemm with ng={ng}, M={M}, N={N}, K={K}, dtype={dtype}, out_dtype={fwd_out_dtype}, main_grad_dtype={main_grad_dtype}, device={device}, TFLOPS={bwd_grad_w_TFLOPS}, dur={bwd_grad_w_dur} ms, bwd_grad_w_efficency={bwd_grad_w_efficency}")
+ else:
+ print(f"-- [bwd_grad_w] Skip skip {bwd_grad_weight_shape_str}")
+
+ all_efficiency[matmul_key]['efficient_factor'] = sum(all_efficiency[matmul_key]['accurate_efficient_factor'].values()) / len(all_efficiency[matmul_key]['accurate_efficient_factor'])
+
+ json.dump(all_efficiency, open(f'{save_root}/grouped_gemm_efficency.json', 'w'), indent=4)
+
+def test_grouped_gemm_efficiency(batch, seq_len, num_experts, ep_size, topk, capacity, hidden_size, moe_intermediate_size, device, MAX_TFLOPS, dtype, grad_reduce_in_bf16, save_root, model):
+ if num_experts % ep_size != 0:
+ return
+ num_local_expert = num_experts // ep_size
+ balance_tokens_per_expert = math.ceil(batch * seq_len * topk / num_experts) * capacity
+ num_tokens_per_local_expert = int(balance_tokens_per_expert * num_experts // num_local_expert)
+ print(f"=== [{model}] Run test for batch={batch}, seq_len={seq_len}, ep_size={ep_size}, num_experts={num_experts}, topk={topk}, capacity={capacity}, hidden_size={hidden_size}, moe_intermediate_size={moe_intermediate_size}, Linear1 device={device}, dtype={dtype}, grad_reduce_in_bf16={grad_reduce_in_bf16}")
+ run_grouped_gemm_fwd_bwd(
+ ng = num_local_expert,
+ M = num_tokens_per_local_expert,
+ K = hidden_size,
+ N = moe_intermediate_size * 2,
+ device=device,
+ dtype=dtype,
+ MAX_TFLOPS=MAX_TFLOPS,
+ grad_reduce_in_bf16 = grad_reduce_in_bf16,
+ save_root = save_root
+ )
+ print(f"=== [{model}] Run test for batch={batch}, seq_len={seq_len}, ep_size={ep_size}, num_experts={num_experts}, topk={topk}, capacity={capacity}, hidden_size={hidden_size}, moe_intermediate_size={moe_intermediate_size}, Linear2 device={device}, dtype={dtype}, grad_reduce_in_bf16={grad_reduce_in_bf16}")
+ run_grouped_gemm_fwd_bwd(
+ ng = num_local_expert,
+ M = num_tokens_per_local_expert,
+ K = moe_intermediate_size,
+ N = hidden_size,
+ device=device,
+ dtype=dtype,
+ MAX_TFLOPS=MAX_TFLOPS,
+ grad_reduce_in_bf16 = grad_reduce_in_bf16,
+ save_root = save_root
+ )
+
+def test(dtype, grad_reduce_in_bf16):
+ system, device, MAX_TFLOPS = get_system_name()
+ save_root = f'{system}_grouped_gemm_efficency'
+ os.makedirs(save_root, exist_ok=True)
+ MODEL_CONFIGS = get_all_test_model_configs()
+ MBS_LIST = get_test_mbs_list()
+ SEQ_LEN_LIST = get_test_seq_len_list()
+ EP_LIST = get_test_ep_list()
+ # tp_list = [1]
+
+ for SEQ_LEN in SEQ_LEN_LIST:
+ for MBS in MBS_LIST:
+ for ep_size in EP_LIST:
+ # for tp_size in tp_list:
+ # tp_size
+ # SEQ_LEN = SEQ_LEN // tp_size
+ for model_config in MODEL_CONFIGS:
+ if model_config.expert_num > 1:
+ test_grouped_gemm_efficiency(batch = MBS,
+ seq_len = SEQ_LEN,
+ num_experts = model_config.expert_num,
+ ep_size = ep_size,
+ topk = model_config.topk,
+ capacity = 1,
+ hidden_size = model_config.hidden_size,
+ moe_intermediate_size = model_config.moe_ffn_hidden_size,
+ device = device,
+ dtype = dtype,
+ MAX_TFLOPS = MAX_TFLOPS,
+ grad_reduce_in_bf16 = grad_reduce_in_bf16,
+ save_root = save_root,
+ model = model_config.model_name)
+if __name__ == "__main__":
+ # test('fp8', False)
+ # test('fp8', True)
+ test('bf16', False)
+ test('bf16', True)
+
+
\ No newline at end of file
diff --git a/simu_tools/efficency_test/utils.py b/simu_tools/efficency_test/utils.py
new file mode 100644
index 0000000..6073155
--- /dev/null
+++ b/simu_tools/efficency_test/utils.py
@@ -0,0 +1,127 @@
+import os
+import time
+import json
+import torch
+
+from simumax.utils import get_simu_model_config, RELEASE_MODELS
+from simumax.core.config import ModelConfig
+
+DEFAULT_MODEL_LIST = ["deepseekv2",
+ "deepseekv3",
+ "deepseek-32b",
+ "deepseek-16b",
+ "deepseek-1b",
+ "llama3-8b",
+ "llama3-70b",
+ "qwen3-32b",
+ "kimi-1T",
+ "mixtral-8x7b"]
+DEFAULT_MBS_LIST = [1, 2, 4]
+DEFAULT_SEQ_LEN_LIST = [4096]
+DEFAULT_TP_LIST = [1, 2, 4, 8]
+DEFAULT_EP_LIST = [1, 2, 4, 8, 16, 64]
+
+PARAM_FILE = os.environ.get("PARAM_FILE", None)
+GLOBAL_PARAMS = None
+if PARAM_FILE:
+ try:
+ GLOBAL_PARAMS = json.load(open(PARAM_FILE))
+ except Exception as e:
+ print(e)
+
+
+def get_test_seq_len_list():
+ if GLOBAL_PARAMS is not None and "seq_len_list" in GLOBAL_PARAMS:
+ return GLOBAL_PARAMS["seq_len_list"]
+ else:
+ return DEFAULT_SEQ_LEN_LIST
+
+def get_test_mbs_list():
+ if GLOBAL_PARAMS is not None and "mbs_list" in GLOBAL_PARAMS:
+ return GLOBAL_PARAMS["mbs_list"]
+ else:
+ return DEFAULT_MBS_LIST
+
+def get_test_ep_list():
+ if GLOBAL_PARAMS is not None and "ep_list" in GLOBAL_PARAMS:
+ return GLOBAL_PARAMS["ep_list"]
+ else:
+ return DEFAULT_EP_LIST
+
+def get_test_tp_list():
+ if GLOBAL_PARAMS is not None and "tp_list" in GLOBAL_PARAMS:
+ return GLOBAL_PARAMS["tp_list"]
+ else:
+ return DEFAULT_TP_LIST
+
+def get_all_test_model_configs():
+ if GLOBAL_PARAMS is not None and "model_list" in GLOBAL_PARAMS:
+ model_list = GLOBAL_PARAMS["model_list"]
+ else:
+ model_list = DEFAULT_MODEL_LIST
+
+ configs = []
+ for model in model_list:
+ if model in RELEASE_MODELS:
+ configs.append(ModelConfig.init_from_config_file(get_simu_model_config(model)))
+ else:
+ print(f"ERROR: Given model_name {model} not found in simumax release model list(path:{RELEASE_MODELS['root']}), skip first!")
+ return configs
+
+def get_torch_profiler(device, profile = False):
+ def trace_handler(prof):
+ rank = 0
+ curr_trace_dir_name = "iteration_" + str(prof.step_num)
+ curr_trace_dir = os.path.join('./profiler_test', curr_trace_dir_name)
+ if not os.path.exists(curr_trace_dir):
+ os.makedirs(curr_trace_dir, exist_ok=True)
+ curr_trace_path = os.path.join(curr_trace_dir, f"rank{rank}.{int(time.time()*1000)}.pt.trace.json")
+ print(f"Dumping profiler traces at step {prof.step_num} to {curr_trace_path}")
+ begin = time.monotonic()
+ prof.export_chrome_trace(curr_trace_path)
+ print(
+ f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
+ )
+ if not profile:
+ return None
+ profiler = torch.profiler.profile(
+ activities=[
+ torch.profiler.ProfilerActivity.CPU,
+ torch.profiler.ProfilerActivity.CUDA,
+ ],
+ schedule=torch.profiler.schedule(wait=2, warmup=3, active=5, skip_first=2,repeat=100),
+ on_trace_ready=trace_handler,
+ record_shapes=True,
+ profile_memory=False,
+ with_stack=True,
+ with_modules=True,
+ )
+ return profiler
+
+def sync_device(device=None):
+ torch.cuda.synchronize()
+
+def get_system_name():
+ system = None
+ device = None
+ MAX_TFLOPS = None
+
+ try:
+ system = torch.cuda.get_device_name(0)
+ device = 'cuda'
+ if 'A100' in system:
+ MAX_TFLOPS = 312
+ elif 'H100' in system:
+ MAX_TFLOPS = 1979
+ except Exception as e:
+ print(f"[ERROR] Info: {e}")
+ exit()
+
+ if os.environ.get('MAX_TFLOPS', None) is not None:
+ MAX_TFLOPS = float(os.environ.get('MAX_TFLOPS'))
+
+ assert MAX_TFLOPS is not None, f"MAX_TFLOPS is None, exit!"
+ assert device is not None, f"device is None, exit!"
+
+ print(f'System: {system}, Device: {device}, Max TFLOPS: {MAX_TFLOPS}')
+ return system, device, MAX_TFLOPS
\ No newline at end of file
diff --git a/simu_tools/megatron_scripts/README-zh.md b/simu_tools/megatron_scripts/README-zh.md
new file mode 100644
index 0000000..64b298d
--- /dev/null
+++ b/simu_tools/megatron_scripts/README-zh.md
@@ -0,0 +1,64 @@
+
+ English|
+ 中文版本
+
+
+
+# 介绍
+SimuMax提供了一套基于Megatron-LM和TE的基准性能测试脚本，可以在英伟达设备上进行llama和deepseek的单机性能测试，用于和SimuMax仿真的性能结果进行校准。
+# 基准性能测试
+## llama3
+支持的测试模型列表：
+- llama2-7b
+- llama3-8b
+- llama3-70b
+
+
+首先在hostfile中添加测试主机ip，然后运行以下命令：
+```bash
+bash run_llama3.sh ${mbs} ${mbc} ${tp_size} ${pp_size} ${model_type} ${test_type} ${num_layers}
+```
+`${test_type}`参数可选`test`和`profile`两个选项，`test`表示基于mock-data运行10个iter的测试训练性能； `profile`表示使用torch.profiler工具导出trace文件， 跑6个iter, 预热1个iter。
+
+例如，运行12层llama3-70b的性能测试，生成的日志文件在./output_llama3_70b中：
+```bash
+bash run_llama3.sh 1 32 1 2 llama3_70b test 12
+```
+
+## deepseek
+支持的测试模型列表：
+- deepseek-v2
+- deepseek-v3
+
+首先在hostfile中添加测试主机ip，然后运行以下命令：
+```bash
+bash run_deepseekv2.sh ${mbs} ${mbc} ${ep_size} ${pp_size} ${model_type} ${test_type} ${num_layers}
+```
+`${test_type}`参数可选`test`和`profile`两个选项，`test`表示基于mock-data运行10个iter的测试训练性能； `profile`表示使用torch.profiler工具导出trace文件， 跑6个iter, 预热1个iter。
+
+例如，运行4层deepseek-v2性能测试，生成的日志文件在./output_deepseek_v2中：
+```bash
+bash run_deepseekv2.sh 1 32 8 1 deepseek_v2 test 4
+```
+运行4层deepseek-v2性能profile，生成的日志文件在./output_deepseek_v2/tboard中：
+```bash
+bash run_deepseekv2.sh 1 32 8 1 deepseek_v2 profile 4
+```
+
+# A100 测试环境说明
+
+框架：
+- te: https://github.com/NVIDIA/TransformerEngine/tree/release_v2.1
+- megatron-lm: https://github.com/NVIDIA/Megatron-LM/releases/tag/v0.12.0rc3\
+
+
+N卡相关硬件/软件环境：
+- NCCL： nvidia-nccl-cu12 2.21.5
+- CUDA 12.6
+- NVIDIA A100 80GB PCIe
+
+A100上使用社区版本flash-mla:
+```bash
+git clone -b head_dim_not_equal https://github.com/defei-coder/flash-attention.git
+cd flash-attention
+MAX_JOBS=8 python setup.py install --verbose > install.log
\ No newline at end of file
diff --git a/simu_tools/megatron_scripts/README.md b/simu_tools/megatron_scripts/README.md
new file mode 100644
index 0000000..0402622
--- /dev/null
+++ b/simu_tools/megatron_scripts/README.md
@@ -0,0 +1,66 @@
+
+ English|
+ 中文版本
+
+
+
+# Introduction
+SimuMax provides a set of benchmark performance testing scripts based on Megatron-LM and Transformer-Engine, which can be used to conduct single-node performance tests for llama and deepseek models on NVIDIA devices. These tests are used to calibrate performance results with SimuMax simulations.
+
+# Benchmark Performance Testing
+## llama3
+Supported test model list:
+- llama2-7b
+- llama3-8b
+- llama3-70b
+
+First, add the test host IP to the hostfile, then run the following command:
+```bash
+bash run_llama3.sh ${mbs} ${mbc} ${tp_size} ${pp_size} ${model_type} ${test_type} ${num_layers}
+```
+
+The ${test_type} parameter has two options: test and profile. test indicates performance testing with mock-data for 10 iterations; profile indicates using torch.profiler to export trace files, running 6 iterations with 1 warm-up iteration.
+
+For example, to run a performance test for a 12-layer llama3-70b, with log files generated in ./output_llama3_70b:
+
+```bash
+bash run_llama3.sh 1 32 1 2 llama3_70b test 12
+```
+
+## deepseek
+Supported test model list:
+- deepseek-v2
+- deepseek-v3
+
+First, add the test host IP to the hostfile, then run the following command:
+```bash
+bash run_deepseekv2.sh ${mbs} ${mbc} ${ep_size} ${pp_size} ${model_type} ${test_type} ${num_layers}
+```
+The ${test_type} parameter has two options: test and profile. test indicates performance testing with mock-data for 10 iterations; profile indicates using torch.profiler to export trace files, running 6 iterations with 1 warm-up iteration.
+
+For example, to run a performance test for a 4-layer deepseek-v2, with log files generated in ./output_deepseek_v2:
+
+```bash
+bash run_deepseekv2.sh 1 32 8 1 deepseek_v2 test 4
+```
+To run a performance profile for a 4-layer deepseek-v2, with log files generated in ./output_deepseek_v2/tboard:
+```bash
+bash run_deepseekv2.sh 1 32 8 1 deepseek_v2 profile 4
+```
+
+# A100 Test Environment
+Framework:
+- TransformerEngine: https://github.com/NVIDIA/TransformerEngine/tree/release_v2.1
+- Megatron-LM: https://github.com/NVIDIA/Megatron-LM/releases/tag/v0.12.0rc3
+
+NVIDIA-related hardware/software environment:
+- NCCL: nvidia-nccl-cu12 2.21.5
+- CUDA 12.6
+- NVIDIA A100 80GB PCIe
+
+Using the community version of flash-mla on A100:
+```bash
+git clone -b head_dim_not_equal https://github.com/defei-coder/flash-attention.git
+cd flash-attention
+MAX_JOBS=8 python setup.py install --verbose > install.log
+```
diff --git a/simu_tools/megatron_scripts/clear.sh b/simu_tools/megatron_scripts/clear.sh
new file mode 100644
index 0000000..d68d32a
--- /dev/null
+++ b/simu_tools/megatron_scripts/clear.sh
@@ -0,0 +1,4 @@
+rm -rf output
+rm -rf logs
+rm -rf checkpoints tboard wandb
+rm -rf ./tmp
\ No newline at end of file
diff --git a/simu_tools/megatron_scripts/hostfile b/simu_tools/megatron_scripts/hostfile
new file mode 100644
index 0000000..7b9ad53
--- /dev/null
+++ b/simu_tools/megatron_scripts/hostfile
@@ -0,0 +1 @@
+127.0.0.1
diff --git a/simu_tools/megatron_scripts/pretrain_deepseekv2.py b/simu_tools/megatron_scripts/pretrain_deepseekv2.py
new file mode 100644
index 0000000..c2a19b4
--- /dev/null
+++ b/simu_tools/megatron_scripts/pretrain_deepseekv2.py
@@ -0,0 +1,308 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+"""Pretrain GPT."""
+
+import os
+import torch
+from functools import partial
+from contextlib import nullcontext
+import inspect
+
+from typing import List, Optional, Tuple, Union
+
+from megatron.training import get_args
+from megatron.training import print_rank_0
+from megatron.training import get_timers
+from megatron.training import get_tokenizer
+from megatron.core import mpu
+from megatron.core.enums import ModelType
+from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
+from megatron.core.datasets.utils import get_blend_from_list
+from megatron.core.datasets.gpt_dataset import GPTDatasetConfig
+from megatron.core.datasets.gpt_dataset import MockGPTDataset, GPTDataset
+from megatron.core.rerun_state_machine import get_rerun_state_machine
+import megatron.legacy.model
+from megatron.core.models.gpt import GPTModel
+from megatron.training import pretrain
+from megatron.core.utils import StragglerDetector
+from megatron.core.transformer.spec_utils import import_module
+from megatron.training.utils import (
+ get_batch_on_this_cp_rank,
+ get_batch_on_this_tp_rank,
+ get_blend_and_blend_per_split,
+)
+from megatron.training.arguments import core_transformer_config_from_args
+from megatron.training.yaml_arguments import core_transformer_config_from_yaml
+from megatron.core.models.gpt.gpt_layer_specs import (
+ get_gpt_decoder_block_spec,
+ get_gpt_layer_local_spec,
+ get_gpt_layer_with_transformer_engine_spec,
+)
+
+
+stimer = StragglerDetector()
+
+def model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megatron.legacy.model.GPTModel]:
+ """Builds the model.
+
+ If you set the use_mcore_models to True, it will return the mcore GPT model and if not the legacy GPT model.
+
+ Args:
+ pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.
+ post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.
+
+
+ Returns:
+ Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model
+ """
+ args = get_args()
+ use_te = args.transformer_impl == "transformer_engine"
+
+ if args.record_memory_history:
+ # torch.cuda.memory._record_memory_history(True,
+ # # keep 100,000 alloc/free events from before the snapshot
+ # trace_alloc_max_entries=100000,
+
+ # # record stack information for the trace events
+ # trace_alloc_record_context=True)
+ torch.cuda.memory._record_memory_history()
+ print_rank_0('building GPT model ...')
+ # Experimental loading arguments from yaml
+ if args.yaml_cfg is not None:
+ config = core_transformer_config_from_yaml(args, "language_model")
+ else:
+ config = core_transformer_config_from_args(args)
+
+ if args.use_legacy_models:
+ model = megatron.legacy.model.GPTModel(
+ config,
+ num_tokentypes=0,
+ parallel_output=True,
+ pre_process=pre_process,
+ post_process=post_process,
+ )
+ else: # using core models
+ if args.spec is not None:
+ transformer_layer_spec = import_module(args.spec)
+ else:
+ if args.num_experts:
+ # Define the decoder block spec
+ transformer_layer_spec = get_gpt_decoder_block_spec(config, use_transformer_engine=use_te)
+ else:
+ # Define the decoder layer spec
+ if use_te:
+ transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+ args.num_experts, args.moe_grouped_gemm,
+ args.qk_layernorm, args.multi_latent_attention, args.moe_use_legacy_grouped_gemm)
+ else:
+ transformer_layer_spec = get_gpt_layer_local_spec(
+ args.num_experts, args.moe_grouped_gemm,
+ args.qk_layernorm, args.multi_latent_attention, args.moe_use_legacy_grouped_gemm)
+
+ build_model_context = nullcontext
+ build_model_context_args = {}
+ if args.fp8_param_gather:
+ try:
+ from transformer_engine.pytorch import fp8_model_init
+
+ build_model_context = fp8_model_init
+ build_model_context_args["enabled"] = True
+ if args.fp8_recipe == 'mxfp8':
+ from transformer_engine.common.recipe import MXFP8BlockScaling, Format
+ recipe = MXFP8BlockScaling()
+ recipe.fp8_format = Format.E4M3
+ build_model_context_args["recipe"] = recipe
+
+ # Check if fp8_model_init supports preserve_high_precision_init_val
+ if "preserve_high_precision_init_val" in inspect.signature(fp8_model_init).parameters:
+ build_model_context_args["preserve_high_precision_init_val"] = True
+ except:
+ raise RuntimeError("--fp8-param-gather requires `fp8_model_init` from TransformerEngine, but not found.")
+
+ with build_model_context(**build_model_context_args):
+ model = GPTModel(
+ config=config,
+ transformer_layer_spec=transformer_layer_spec,
+ vocab_size=args.padded_vocab_size,
+ max_sequence_length=args.max_position_embeddings,
+ pre_process=pre_process,
+ post_process=post_process,
+ fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+ parallel_output=True,
+ share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+ position_embedding_type=args.position_embedding_type,
+ rotary_percent=args.rotary_percent,
+ rotary_base=args.rotary_base,
+ rope_scaling=args.use_rope_scaling
+ )
+
+ return model
+
+
+def get_batch(data_iterator):
+ """Generate a batch."""
+
+ # TODO: this is pretty hacky, find a better way
+ if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()):
+ return None, None, None, None, None
+
+ # get batches based on the TP rank you are on
+ batch = get_batch_on_this_tp_rank(data_iterator)
+
+ # slice batch along sequence dimension for context parallelism
+ batch = get_batch_on_this_cp_rank(batch)
+
+ return batch.values()
+
+
+def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
+ """Loss function.
+
+ Args:
+ loss_mask (torch.Tensor): Used to mask out some portions of the loss
+ output_tensor (torch.Tensor): The tensor with the losses
+
+ Returns:
+ the loss scalar for this micro-batch
+ the number of non-padded tokens in this microbatch
+ a dict containing reporting metrics on the loss and number of tokens across
+ the data parallel ranks
+ """
+ args = get_args()
+
+ losses = output_tensor.float()
+ loss_mask = loss_mask.view(-1).float()
+ total_tokens = loss_mask.sum()
+ loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), total_tokens.view(1)])
+
+ if args.context_parallel_size > 1:
+ torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group())
+
+ # Check individual rank losses are not NaN prior to DP all-reduce.
+ rerun_state_machine = get_rerun_state_machine()
+ if args.check_for_nan_in_loss_and_grad:
+ rerun_state_machine.validate_result(
+ result=loss[0],
+ rejection_func=torch.isnan,
+ message="found NaN in local forward loss calculation",
+ tolerance=0.0, # forward pass calculations are determinisic
+ fatal=True,
+ )
+ # Check for spiky loss
+ if args.check_for_spiky_loss:
+ rerun_state_machine.validate_result(
+ result=loss[0],
+ rejection_func=partial(rerun_state_machine.is_spiky_loss, threshold=SPIKY_LOSS_PERC),
+ message="Spiky loss",
+ tolerance=0.0, # forward pass calculations are determinisic
+ fatal=False,
+ )
+ # Reduce loss for logging.
+ reporting_loss = loss.clone().detach()
+ if not int(os.getenv("NO_LOSS_REDUCE", 0)): #TODO:(huang.huang) will influence the loss reported Now!
+ torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group())
+
+
+ local_num_tokens = loss[1].clone().detach().to(torch.int)
+ return (
+ loss[0] * args.context_parallel_size,
+ local_num_tokens,
+ {'lm loss': (reporting_loss[0], reporting_loss[1])},
+ )
+
+
+def forward_step(data_iterator, model: GPTModel):
+ """Forward training step.
+
+ Args:
+ data_iterator : Input data iterator
+ model (GPTModel): The GPT Model
+ """
+ args = get_args()
+ timers = get_timers()
+
+ # Get the batch.
+ timers('batch-generator', log_level=2).start()
+ global stimer
+ with stimer(bdata=True):
+ tokens, labels, loss_mask, attention_mask, position_ids = get_batch(
+ data_iterator)
+ timers('batch-generator').stop()
+
+ with stimer:
+ output_tensor = model(tokens, position_ids, attention_mask,
+ labels=labels)
+
+ return output_tensor, partial(loss_func, loss_mask)
+
+
+def is_dataset_built_on_rank():
+ return (
+ mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage()
+ ) and mpu.get_tensor_model_parallel_rank() == 0
+
+
+def core_gpt_dataset_config_from_args(args):
+ tokenizer = get_tokenizer()
+
+ return GPTDatasetConfig(
+ random_seed=args.seed,
+ sequence_length=args.seq_length,
+ blend=get_blend_from_list(args.data_path),
+ blend_per_split=[
+ get_blend_from_list(args.train_data_path),
+ get_blend_from_list(args.valid_data_path),
+ get_blend_from_list(args.test_data_path)
+ ],
+ split=args.split,
+ num_dataset_builder_threads=args.num_dataset_builder_threads,
+ path_to_cache=args.data_cache_path,
+ mmap_bin_files=args.mmap_bin_files,
+ tokenizer=tokenizer,
+ reset_position_ids=args.reset_position_ids,
+ reset_attention_mask=args.reset_attention_mask,
+ eod_mask_loss=args.eod_mask_loss,
+ create_attention_mask=args.create_attention_mask_in_dataloader,
+ )
+
+
+def train_valid_test_datasets_provider(train_val_test_num_samples):
+ """Build the train test and validation datasets.
+
+ Args:
+ train_val_test_num_samples : A list containing the number of samples in train test and validation.
+ """
+ args = get_args()
+
+ config = core_gpt_dataset_config_from_args(args)
+
+ if args.mock_data:
+ dataset_type = MockGPTDataset
+ else:
+ dataset_type = GPTDataset
+
+ print_rank_0("> building train, validation, and test datasets for GPT ...")
+
+ train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
+ dataset_type,
+ train_val_test_num_samples,
+ is_dataset_built_on_rank,
+ config
+ ).build()
+
+ print_rank_0("> finished creating GPT datasets ...")
+
+ return train_ds, valid_ds, test_ds
+
+
+if __name__ == "__main__":
+
+ # Temporary for transition to core datasets
+ train_valid_test_datasets_provider.is_distributed = True
+
+ pretrain(
+ train_valid_test_datasets_provider,
+ model_provider,
+ ModelType.encoder_or_decoder,
+ forward_step,
+ args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
+ )
diff --git a/simu_tools/megatron_scripts/pretrain_gpt.py b/simu_tools/megatron_scripts/pretrain_gpt.py
new file mode 100644
index 0000000..cbc0ad9
--- /dev/null
+++ b/simu_tools/megatron_scripts/pretrain_gpt.py
@@ -0,0 +1,280 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+"""Pretrain GPT."""
+
+import os
+import torch
+from functools import partial
+from contextlib import nullcontext
+import inspect
+
+from typing import Union
+
+from megatron.training import get_args
+from megatron.training import print_rank_0
+from megatron.training import get_timers
+from megatron.training import get_tokenizer
+from megatron.core import mpu
+from megatron.core.enums import ModelType
+from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
+from megatron.core.datasets.utils import get_blend_from_list
+from megatron.core.datasets.gpt_dataset import GPTDatasetConfig
+from megatron.core.datasets.gpt_dataset import MockGPTDataset, GPTDataset
+import megatron.legacy.model
+from megatron.core.models.gpt import GPTModel
+from megatron.training import pretrain
+from megatron.core.utils import StragglerDetector
+from megatron.core.transformer.spec_utils import import_module
+from megatron.training.utils import (
+ get_batch_on_this_cp_rank,
+ get_batch_on_this_tp_rank,
+)
+from megatron.training.arguments import core_transformer_config_from_args
+from megatron.training.yaml_arguments import core_transformer_config_from_yaml
+from megatron.core.models.gpt.gpt_layer_specs import (
+ get_gpt_layer_local_spec,
+ get_gpt_layer_with_transformer_engine_spec,
+)
+
+
+stimer = StragglerDetector()
+
+def model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megatron.legacy.model.GPTModel]:
+ """Builds the model.
+
+ If you set the use_legacy_models to True, it will return the legacy GPT model and if not the mcore GPT model.
+
+ Args:
+ pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.
+ post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.
+
+
+ Returns:
+ Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model
+ """
+ args = get_args()
+ use_te = args.transformer_impl == "transformer_engine"
+
+ print_rank_0('building GPT model ...')
+ # Experimental loading arguments from yaml
+ if args.yaml_cfg is not None:
+ config = core_transformer_config_from_yaml(args, "language_model")
+ else:
+ config = core_transformer_config_from_args(args)
+
+ if args.use_legacy_models:
+ model = megatron.legacy.model.GPTModel(
+ config,
+ num_tokentypes=0,
+ parallel_output=True,
+ pre_process=pre_process,
+ post_process=post_process,
+ )
+ else: # using core models
+ if args.spec is not None:
+ transformer_layer_spec = import_module(args.spec)
+ else:
+ if use_te:
+ transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+ args.num_experts,
+ args.moe_grouped_gemm,
+ args.qk_layernorm,
+ args.multi_latent_attention,
+ args.moe_use_legacy_grouped_gemm
+ )
+ else:
+ transformer_layer_spec = get_gpt_layer_local_spec(
+ args.num_experts,
+ args.moe_grouped_gemm,
+ args.qk_layernorm
+ )
+
+ build_model_context = nullcontext
+ build_model_context_args = {}
+ if args.fp8_param_gather:
+ try:
+ from transformer_engine.pytorch import fp8_model_init
+
+ build_model_context = fp8_model_init
+ build_model_context_args["enabled"] = True
+
+ # Check if fp8_model_init supports preserve_high_precision_init_val
+ if "preserve_high_precision_init_val" in inspect.signature(fp8_model_init).parameters:
+ build_model_context_args["preserve_high_precision_init_val"] = True
+ except:
+ raise RuntimeError("--fp8-param-gather requires `fp8_model_init` from TransformerEngine, but not found.")
+
+ with build_model_context(**build_model_context_args):
+ model = GPTModel(
+ config=config,
+ transformer_layer_spec=transformer_layer_spec,
+ vocab_size=args.padded_vocab_size,
+ max_sequence_length=args.max_position_embeddings,
+ pre_process=pre_process,
+ post_process=post_process,
+ fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+ parallel_output=True,
+ share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+ position_embedding_type=args.position_embedding_type,
+ rotary_percent=args.rotary_percent,
+ rotary_base=args.rotary_base
+ )
+
+ return model
+
+
+def get_batch(data_iterator):
+ """Generate a batch."""
+
+ # TODO: this is pretty hacky, find a better way
+ if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()):
+ return None, None, None, None, None
+
+ # get batches based on the TP rank you are on
+ batch = get_batch_on_this_tp_rank(data_iterator)
+
+ # slice batch along sequence dimension for context parallelism
+ batch = get_batch_on_this_cp_rank(batch)
+
+ return batch.values()
+
+
+def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
+ """Loss function.
+
+ Args:
+ loss_mask (torch.Tensor): Used to mask out some portions of the loss
+ output_tensor (torch.Tensor): The tensor with the losses
+
+ Returns:
+ the loss scalar for this micro-batch
+ the number of non-padded tokens in this microbatch
+ a dict containing reporting metrics on the loss and number of tokens across
+ the data parallel ranks
+ """
+ args = get_args()
+
+ losses = output_tensor.float()
+ loss_mask = loss_mask.view(-1).float()
+ total_tokens = loss_mask.sum()
+ loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), total_tokens.view(1)])
+
+ if args.context_parallel_size > 1:
+ torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group())
+
+ # Check individual rank losses are not NaN prior to DP all-reduce.
+ if args.check_for_nan_in_loss_and_grad:
+ global_rank = torch.distributed.get_rank()
+ assert not loss[0].isnan(), (
+ f'Rank {global_rank}: found NaN in local forward loss calculation. '
+ f'Device: {torch.cuda.current_device()}, node: {os.uname()[1]}'
+ )
+
+ # Reduce loss for logging.
+ reporting_loss = loss.clone().detach()
+ if not int(os.getenv("NO_LOSS_REDUCE", 0)): #TODO:(huang.huang) will influence the loss reported Now!
+ torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group())
+ local_num_tokens = loss[1].clone().detach().to(torch.int)
+ return (
+ loss[0] * args.context_parallel_size,
+ local_num_tokens,
+ {'lm loss': (reporting_loss[0], reporting_loss[1])},
+ )
+
+
+def forward_step(data_iterator, model: GPTModel):
+ """Forward training step.
+
+ Args:
+ data_iterator : Input data iterator
+ model (GPTModel): The GPT Model
+ """
+ args = get_args()
+ timers = get_timers()
+
+ # Get the batch.
+ timers('batch-generator', log_level=2).start()
+ global stimer
+ with stimer(bdata=True):
+ tokens, labels, loss_mask, attention_mask, position_ids = get_batch(
+ data_iterator)
+ timers('batch-generator').stop()
+
+ with stimer:
+ output_tensor = model(tokens, position_ids, attention_mask,
+ labels=labels)
+
+ return output_tensor, partial(loss_func, loss_mask)
+
+
+def is_dataset_built_on_rank():
+ return (
+ mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage()
+ ) and mpu.get_tensor_model_parallel_rank() == 0
+
+
+def core_gpt_dataset_config_from_args(args):
+ tokenizer = get_tokenizer()
+
+ return GPTDatasetConfig(
+ random_seed=args.seed,
+ sequence_length=args.seq_length,
+ blend=get_blend_from_list(args.data_path),
+ blend_per_split=[
+ get_blend_from_list(args.train_data_path),
+ get_blend_from_list(args.valid_data_path),
+ get_blend_from_list(args.test_data_path)
+ ],
+ # renormalize_blend_weights=args.renormalize_blend_weights,
+ split=args.split,
+ num_dataset_builder_threads=args.num_dataset_builder_threads,
+ path_to_cache=args.data_cache_path,
+ mmap_bin_files=args.mmap_bin_files,
+ tokenizer=tokenizer,
+ reset_position_ids=args.reset_position_ids,
+ reset_attention_mask=args.reset_attention_mask,
+ eod_mask_loss=args.eod_mask_loss,
+ create_attention_mask=args.create_attention_mask_in_dataloader,
+ )
+
+
+def train_valid_test_datasets_provider(train_val_test_num_samples):
+ """Build the train test and validation datasets.
+
+ Args:
+ train_val_test_num_samples : A list containing the number of samples in train test and validation.
+ """
+ args = get_args()
+
+ config = core_gpt_dataset_config_from_args(args)
+
+ if args.mock_data:
+ dataset_type = MockGPTDataset
+ else:
+ dataset_type = GPTDataset
+
+ print_rank_0("> building train, validation, and test datasets for GPT ...")
+
+ train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
+ dataset_type,
+ train_val_test_num_samples,
+ is_dataset_built_on_rank,
+ config
+ ).build()
+
+ print_rank_0("> finished creating GPT datasets ...")
+
+ return train_ds, valid_ds, test_ds
+
+
+if __name__ == "__main__":
+
+ # Temporary for transition to core datasets
+ train_valid_test_datasets_provider.is_distributed = True
+
+ pretrain(
+ train_valid_test_datasets_provider,
+ model_provider,
+ ModelType.encoder_or_decoder,
+ forward_step,
+ args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
+ )
diff --git a/simu_tools/megatron_scripts/run.sh b/simu_tools/megatron_scripts/run.sh
new file mode 100644
index 0000000..d69c43d
--- /dev/null
+++ b/simu_tools/megatron_scripts/run.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+mbc_list=(4 8 32)
+
+# 遍历每个 mbc 值
+for mbc in "${mbc_list[@]}"; do
+ bash run_deepseekv2.sh 1 $mbc 8 1 deepseek_v2 test 4
+ bash run_deepseekv2.sh 1 $mbc 4 2 deepseek_v2 test 4
+ bash run_llama3.sh 1 $mbc 1 2 llama3_8b test 32
+ bash run_llama3.sh 1 $mbc 2 1 llama3_8b test 32
+ bash run_llama3.sh 1 $mbc 4 1 llama3_8b test 32
+ bash run_llama3.sh 1 $mbc 8 1 llama3_8b test 32
+
+ bash run_llama3.sh 1 $mbc 1 2 llama3_70b test 12
+ bash run_llama3.sh 1 $mbc 2 1 llama3_70b test 12
+ bash run_llama3.sh 1 $mbc 4 1 llama3_70b test 12
+ bash run_llama3.sh 1 $mbc 8 1 llama3_70b test 12
+
+ # bash run_deepseekv2.sh 1 $mbc 8 1 deepseek_v2 test 4 fp8
+done
diff --git a/simu_tools/megatron_scripts/run_deepseekv2.sh b/simu_tools/megatron_scripts/run_deepseekv2.sh
new file mode 100755
index 0000000..acaf697
--- /dev/null
+++ b/simu_tools/megatron_scripts/run_deepseekv2.sh
@@ -0,0 +1,83 @@
+#!/bin/bash
+LOCK_FILE="/tmp/$(basename "$0").lock"
+exec 200>"$LOCK_FILE"
+flock -n 200 || {
+ echo "等待前一个实例完成..."
+ flock 200
+ echo "开始执行..."
+}
+
+tp_size=1
+world_size=8
+micro_batch_size=$1
+num_microbatches=$2
+ep_size=$3
+pp_size=$4
+model_type=$5
+test_type=$6
+num_layer=$7
+dtype=${8:-"bf16"}
+(( dp_size = $world_size / ($tp_size * $pp_size) ))
+(( global_batch_size = $micro_batch_size * $num_microbatches * $dp_size ))
+
+# echo -e "\033[32mdp_size: $dp_size, pp_size: $pp_size, ep_size: $ep_size, \
+# world_size: $world_size, micro_batch_size: $micro_batch_size, \
+# num_microbatches: $num_microbatches, global_batch_size: $global_batch_size, model_type=$model_type, est_type=$test_type, num_layer=$num_layer, dtype=$dtype\033[0m"
+current_time=$(date "+%Y%m%d_%H%M%S")
+output_dir="./output_${model_type}/$current_time"
+mkdir -p ${output_dir}
+
+set -u
+ work_home=$(realpath ./)
+ patch_home=""
+ megatron_home="./Megatron-LM"
+ example="tp${tp_size}_pp${pp_size}_dp${dp_size}_mbs${micro_batch_size}_mbc${num_microbatches}_gbs${global_batch_size}_gpus${world_size}"
+ data_path=${data_path:-"/home/dist/yehua/llama2_dataset/llama_00_text_document"}
+ hostfile=./hostfile
+ log_file=${output_dir}/$example.log
+ script_file=./scripts/run_pretrain_deepseekv2.sh
+ rdzv_id=$current_time
+ master_port=12345
+set +u
+
+cmd_env=(
+ "WORK_HOME=${work_home}"
+ "PATCH_HOME=${patch_home}"
+ "MEGATRON_HOME=${megatron_home}"
+ "EXAMPLE=${example}"
+ "HOSTFILE=${hostfile}"
+ "TP_SIZE=${tp_size}"
+ "PP_SIZE=${pp_size}"
+ "EP_SIZE=${ep_size}"
+ "MICRO_BATCH_SIZE=${micro_batch_size}"
+ "GLOBAL_BATCH_SIZE=${global_batch_size}"
+ "RDZV_ID=${rdzv_id}"
+ "MASTER_PORT=${master_port}"
+ "TEST_TYPE=${test_type}"
+ "NUM_LAYER=${num_layer}"
+ "MODEL_TYPE=${model_type}"
+ "DTYPE=${dtype}"
+ "OUTPUT_DIR=${output_dir}"
+)
+
+# 执行
+cmd="cd ${work_home} && ${cmd_env[@]} bash ${script_file}"
+
+hostlist=$(grep -v '^#\|^$' $hostfile | awk '{print $1}' | xargs)
+
+# Check if hostlist is empty
+if [ -z "$hostlist" ]; then
+ echo "Error: hostlist is empty. Please add IP addresses to the hostfile."
+ exit 1
+fi
+
+for host in ${hostlist[@]}; do
+ # cmd_ssh="$cmd > ${output_dir}/$host.log 2>&1"
+ cmd_ssh="$cmd > ${output_dir}/$host.log"
+ echo $cmd_ssh
+ ssh -n $host $cmd_ssh &
+ # eval $cmd_ssh
+done
+
+wait
+bash stop_all.sh
\ No newline at end of file
diff --git a/simu_tools/megatron_scripts/run_llama3.sh b/simu_tools/megatron_scripts/run_llama3.sh
new file mode 100644
index 0000000..85ac80e
--- /dev/null
+++ b/simu_tools/megatron_scripts/run_llama3.sh
@@ -0,0 +1,78 @@
+#!/bin/bash
+LOCK_FILE="/tmp/megatron_diff.lock"
+exec 200>"$LOCK_FILE"
+flock -n 200 || {
+ echo "等待前一个实例完成..."
+ flock 200
+ echo "开始执行..."
+}
+
+micro_batch_size=$1
+num_microbatches=$2
+tp_size=$3
+pp_size=$4
+model_type=$5
+test_type=$6
+num_layer=$7
+dtype=${8:-"bf16"}
+world_size=8
+# echo "pp_size=$pp_size, tp_size=$tp_size, micro_batch_size=$micro_batch_size, num_microbatches=$num_microbatches, model_type=$model_type, test_type=$test_type, num_layer=$num_layer"
+# echo -e "\033[32mdp_size: $dp_size, pp_size: $pp_size, ep_size: $ep_size, \
+# world_size: $world_size, micro_batch_size: $micro_batch_size, \
+# num_microbatches: $num_microbatches, global_batch_size: $global_batch_size, model_type=$model_type, test_type=$test_type, num_layer=$num_layer, dtype=$dtype\033[0m"
+
+(( dp_size = $world_size / ($tp_size * $pp_size) ))
+(( global_batch_size = $micro_batch_size * $num_microbatches * $dp_size ))
+current_time=$(date "+%Y%m%d_%H%M%S")
+output_dir="./output_${model_type}/$current_time"
+mkdir -p ${output_dir}
+
+set -u
+ work_home=$(realpath ./)
+ patch_home=""
+ megatron_home="./Megatron-LM"
+ example="tp${tp_size}_pp${pp_size}_dp${dp_size}_mbs${micro_batch_size}_mbc${num_microbatches}_gbs${global_batch_size}_gpus${world_size}"
+ hostfile=./hostfile
+ log_file=${output_dir}/$example.log
+ script_file=./scripts/run_pretrain_llama3.sh
+ rdzv_id=$current_time
+set +u
+
+cmd_env=(
+ "WORK_HOME=${work_home}"
+ "PATCH_HOME=${patch_home}"
+ "MEGATRON_HOME=${megatron_home}"
+ "EXAMPLE=${example}" # 注意：这里从 EXAMPLE 改为了 EXPNAME
+ "HOSTFILE=${hostfile}"
+ "TP_SIZE=${tp_size}"
+ "PP_SIZE=${pp_size}"
+ "MICRO_BATCH_SIZE=${micro_batch_size}"
+ "GLOBAL_BATCH_SIZE=${global_batch_size}"
+ "RDZV_ID=${rdzv_id}"
+ "MODEL_TYPE=${model_type}"
+ "TEST_TYPE=${test_type}"
+ "NUM_LAYERS=${num_layer}" # 注意：这里从 NUM_LAYER 改为了 NUM_LAYERS
+ "DTYPE=${dtype}"
+ "OUTPUT_DIR=${output_dir}"
+)
+
+# 执行
+cmd="cd ${work_home} && ${cmd_env[@]} bash ${script_file}"
+hostlist=$(grep -v '^#\|^$' $hostfile | awk '{print $1}' | xargs)
+
+# Check if hostlist is empty
+if [ -z "$hostlist" ]; then
+ echo "Error: hostlist is empty. Please add IP addresses to the hostfile."
+ exit 1
+fi
+
+for host in ${hostlist[@]}; do
+ # cmd_ssh="$cmd > ${output_dir}/$host.log 2>&1"
+ cmd_ssh="$cmd > ${output_dir}/$host.log"
+ echo $cmd_ssh
+ ssh -n $host $cmd_ssh &
+ # eval $cmd_ssh
+done
+
+wait
+bash stop_all.sh
\ No newline at end of file
diff --git a/simu_tools/megatron_scripts/scripts/run_pretrain_deepseekv2.sh b/simu_tools/megatron_scripts/scripts/run_pretrain_deepseekv2.sh
new file mode 100644
index 0000000..26416cd
--- /dev/null
+++ b/simu_tools/megatron_scripts/scripts/run_pretrain_deepseekv2.sh
@@ -0,0 +1,315 @@
+#!/bin/bash
+
+set -u
+WORK_HOME=${WORK_HOME:-"./"}
+PATCH_HOME=${PATCH_HOME:-"./"}
+MEGATRON_HOME=${MEGATRON_HOME:-"./"}
+EXAMPLE=${EXAMPLE:-"test"}
+HOSTFILE=${HOSTFILE:-"./hostfile"}
+DATA_DIR=${DATA_DIR:-"./"}
+TP_SIZE=${TP_SIZE:-1}
+PP_SIZE=${PP_SIZE:-1}
+EP_SIZE=${EP_SIZE:-1}
+MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
+GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-1}
+TOKENIZED_MODEL=${TOKENIZED_MODEL:-""}
+RDZV_ID=${RDZV_ID:-"0"}
+MASTER_PORT=${MASTER_PORT:-"12345"}
+TEST_TYPE=${TEST_TYPE:-"test"}
+NUM_LAYER=${NUM_LAYER:-"1"}
+MODEL_TYPE=${MODEL_TYPE:-""}
+DTYPE=${DTYPE:-"bf16"}
+SCRIPT_FILE=${SCRIPT_FILE:-""}
+OUTPUT_DIR=${OUTPUT_DIR:-"./output_${MODEL_TYPE}"}
+set +u
+
+
+if [ "${TEST_TYPE}" == "profile" ]; then
+ PROFILER_SAVE_PATH=${WORK_HOME}/profiler_result_ds_236b_${EXAMPLE}
+ rm -rf ${PROFILER_SAVE_PATH}
+ echo "Profiler save path: ${PROFILER_SAVE_PATH}"
+ # export ENABLE_PROFILER=1
+ # export PROFILER_FREQ=4
+ # export PROFILER_SAVE_DIR=${PROFILER_SAVE_PATH}
+ PROFILE_ARGS=(
+ --profile
+ --profile-step-start 4
+ --profile-step-end 6
+ --use-pytorch-profiler
+ )
+ TRAIN_ITERS=6
+ WARMUP_ITERS=1
+else
+ PROFILE_ARGS=()
+ TRAIN_ITERS=10
+ WARMUP_ITERS=5
+fi
+
+export OMP_NUM_THREADS=4
+export CUDA_VISIBLE_DEVICES='0,1,2,3,4,5,6,7'
+export ACCELERATOR_BACKEND="cuda"
+
+
+export MOE_ROUTER_GROUP_TOPK=${MOE_ROUTER_GROUP_TOPK:-3}
+
+export ENABLE_ZERO_BUBBLE=0 # if set 1, Enable zero_bubble
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export CPU_OPTIMIZER_PRECISION_AWARE_RECONFIG=${CPU_OPTIMIZER_PRECISION_AWARE_RECONFIG:-0}
+
+export ENABLE_D2H_IN_PERMUTATION=0
+# export NO_LOSS_REDUCE=1
+
+export PYTHONPATH=${MEGATRON_HOME}:${PATCH_HOME}:$PYTHONPATH
+echo $PYTHONPATH
+
+
+# CHECKPOINT_PATH=$OUTPUT_DIR/checkpoints/$EXAMPLE
+# mkdir -p $CHECKPOINT_PATH
+DATA_PATH=$DATA_DIR
+
+
+# LOG_PATH=$OUTPUT_DIR/logs/$EXAMPLE
+# mkdir -p $LOG_PATH
+# cp $0 $LOG_PATH/
+TB_PATH=$OUTPUT_DIR/tboard/$EXAMPLE
+mkdir -p $TB_PATH
+# WB_PATH=$OUTPUT_DIR/wandb/$EXAMPLE
+# mkdir -p $WB_PATH
+
+
+export NODE_ADDR=$(ip a|grep inet|grep -v 127.0.0.1|grep -v inet6|awk '{print $2;}'|tr -d "addr:"|head -n1 | cut -d '/' -f1) # tail for cuda
+export GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+export NUM_NODES=$(cat $HOSTFILE | wc -l)
+export MASTER_ADDR=$(head -n1 $HOSTFILE | awk '{print $1;}')
+export NODE_RANK=$(awk -v node_addr="$NODE_ADDR" '{ranks[$1]=(FNR-1);} END {print ranks[node_addr];}' $HOSTFILE)
+export MASTER_PORT=${MASTER_PORT:-12356}
+
+LOG_DIR=${OUTPUT_DIR}/${DTYPE}_$EXAMPLE
+echo "Distributed log_dir: $LOG_DIR"
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --node_rank $NODE_RANK
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+ --log_dir $LOG_DIR
+ --redirects ${LOG_REDIRECTS_LEVEL:-3}
+)
+echo $GPUS_PER_NODE $NUM_NODES $NODE_RANK $MASTER_ADDR $MASTER_PORT
+
+NUM_LAYERS_MINUS_ONE=$((NUM_LAYER - 1))
+NUM_LAYERS_MINUS_THREE=$((NUM_LAYER - 3))
+if [ "$MODEL_TYPE" == "deepseek_v2" ]; then
+ HEAD_NUM=128
+ # KV_HEAD_NUM=128 # not used
+ HIDDEN_SIZE=5120
+ export MOE_NUM_EXPERTS=160
+ MOE_LAYER_FREQ="([0]*1+[1]*${NUM_LAYERS_MINUS_ONE})*1"
+ MOE_FFN_HIDDEN_SIZE=1536
+ FFN_HIDDEN_SIZE=12288
+ VOCAB_SIZE=102400
+ TOPK=6
+ MOE_SHARED_FFN_HIDDEN_SIZE=3072
+ echo "MOE_LAYER_FREQ: $MOE_LAYER_FREQ"
+
+elif [ "$MODEL_TYPE" == "deepseek_v3" ]; then
+ HEAD_NUM=128
+ # KV_HEAD_NUM=128 # not used
+ HIDDEN_SIZE=7168
+ export MOE_NUM_EXPERTS=256
+ MOE_LAYER_FREQ="([0]*1+[1]*${NUM_LAYERS_MINUS_ONE})*1"
+ MOE_FFN_HIDDEN_SIZE=2048
+ FFN_HIDDEN_SIZE=18432
+ VOCAB_SIZE=129280
+ TOPK=8
+ MOE_SHARED_FFN_HIDDEN_SIZE=2048
+
+elif [ "$MODEL_TYPE" == "kimi_1t" ]; then
+ HEAD_NUM=64
+ # KV_HEAD_NUM=64 # not used
+ HIDDEN_SIZE=7168
+ export MOE_NUM_EXPERTS=384
+ MOE_LAYER_FREQ="([0]*1+[1]*${NUM_LAYERS_MINUS_ONE})*1"
+ MOE_FFN_HIDDEN_SIZE=2048
+ FFN_HIDDEN_SIZE=18432
+ VOCAB_SIZE=163840
+ TOPK=8
+ MOE_SHARED_FFN_HIDDEN_SIZE=2048
+else echo "Unknown model type: $MODEL_TYPE"
+fi
+MODEL_ARGS=(
+ --num-layers ${NUM_LAYER} # 60 ds 236b:ep8 pp16
+ --hidden-size $HIDDEN_SIZE # dsv2 = 5120, v3=7168
+ --num-attention-heads $HEAD_NUM # dsv2/3=128, kimi-1T=64
+ --seq-length 4096
+ --max-position-embeddings 4096
+ --norm-epsilon 1e-6
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --disable-bias-linear
+ --vocab-size $VOCAB_SIZE # dsv2 = 102400, v3=129280
+ --ffn-hidden-size $FFN_HIDDEN_SIZE # 12288 for dense, but for sequentialMLP, moe-ffn-hidden-size=1536 # dsv2 = 12288, v3=18432
+ --position-embedding-type rope
+ --no-position-embedding
+ --swiglu
+ --normalization RMSNorm
+ --untie-embeddings-and-output-weights
+ # --cross-entropy-loss-fusion
+
+)
+
+
+MOE_ARGS=(
+ --num-experts ${MOE_NUM_EXPERTS}
+ --expert-model-parallel-size $EP_SIZE
+ --moe-token-dispatcher-type alltoall
+ --moe-router-score-function softmax
+ --moe-router-num-groups $EP_SIZE
+ --moe-router-group-topk ${MOE_ROUTER_GROUP_TOPK}
+ --moe-router-load-balancing-type seq_aux_loss
+ --moe-router-topk $TOPK
+ --moe-router-pre-softmax #deepseek use pre-softmax
+ --moe-router-topk-scaling-factor 16 # pre-softmax need scaling
+ --moe-aux-loss-coeff 3e-3
+ --moe-expert-capacity-factor 1
+
+ --moe-ffn-hidden-size $MOE_FFN_HIDDEN_SIZE # ds_v2 = 1536, v3 = 2048
+ --moe-shared-expert-intermediate-size $MOE_SHARED_FFN_HIDDEN_SIZE # ds_v2 = 3072, v3= 2048
+ --moe-layer-freq "$MOE_LAYER_FREQ"
+ --moe-grouped-gemm
+ --moe-permute-fusion
+ --moe-pad-expert-input-to-capacity
+)
+
+# 24414062 1T
+TRAINING_ARGS=(
+ --seed 42
+ --micro-batch-size $MICRO_BATCH_SIZE
+ --global-batch-size $GLOBAL_BATCH_SIZE
+ # --train-samples 24414062
+ --init-method-std 0.006 # 0.02 in HF config, but 0.006 in the paper
+ --use-mcore-models
+ # --no-gradient-accumulation-fusion
+ --no-bias-dropout-fusion
+ # --no-bias-swiglu-fusion
+ --use-distributed-optimizer
+ --use-flash-attn
+ --sequence-parallel
+ --recompute-granularity full
+ --recompute-method block
+ --recompute-num-layers 0
+ --distributed-backend nccl
+ --multi-latent-attention
+ --qk-layernorm
+ --mock-data
+)
+
+MLA_ARGS=(
+ --q-lora-rank 1536
+ --kv-lora-rank 512
+ --qk-head-dim 128
+ --qk-pos-emb-head-dim 64
+ --v-head-dim 128
+ --rotary-scaling-factor 1
+)
+
+REGULARIZATION_ARGS=(
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.95
+ --clip-grad 1.0
+)
+
+WARMUP_STEPS=2000
+WARMUP_SAMPLES=$((WARMUP_STEPS * GLOBAL_BATCH_SIZE))
+
+LEARNING_RATE_ARGS=(
+ # --lr 1.5e-5
+ # --lr-decay-style cosine
+ # --lr-warmup-samples ${WARMUP_SAMPLES}
+ # --min-lr 1.5e-6
+ --initial-loss-scale 65536
+ --min-loss-scale 1.0
+
+ # 下面是ds 32b参数
+ --train-iters ${TRAIN_ITERS}
+ --lr-warmup-iters ${WARMUP_ITERS}
+ --lr 2.4e-4
+ --lr-decay-style cosine
+ --min-lr 1e-6
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size $TP_SIZE
+ --pipeline-model-parallel-size $PP_SIZE
+)
+
+MIXED_PRECISION_ARGS=(
+ --bf16
+ --attention-softmax-in-fp32
+ --no-masked-softmax-fusion
+ --accumulate-allreduce-grads-in-fp32
+)
+
+DATA_ARGS=(
+ # --data-path $DATA_PATH
+ --tokenizer-type NullTokenizer
+ # --tokenizer-model ${TOKENIZED_MODEL}
+ --split 1
+ #--dataloader-type mtepx #default single
+)
+
+EVAL_AND_LOGGING_ARGS=(
+ --log-interval 1
+ --log-throughput
+ --save-interval 100000
+ --eval-interval 1
+ # --save $CHECKPOINT_PATH
+ # --load $CHECKPOINT_PATH
+ --eval-iters 0
+ --tensorboard-dir $TB_PATH
+)
+
+
+if [ "$DTYPE" == "bf16" ]; then
+ TRANSFORMER_ENGINE_ARGS=(
+ --transformer-impl transformer_engine
+ )
+else
+ TRANSFORMER_ENGINE_ARGS=(
+ --transformer-impl transformer_engine
+ --fp8-format hybrid
+ --fp8-param-gather
+ )
+fi
+
+
+cmd="torchrun ${DISTRIBUTED_ARGS[@]} ./pretrain_deepseekv2.py \
+ ${MODEL_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${REGULARIZATION_ARGS[@]}
+ ${LEARNING_RATE_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${MIXED_PRECISION_ARGS[@]}
+ ${DATA_ARGS[@]} \
+ ${MOE_ARGS[@]} \
+ ${MLA_ARGS[@]} \
+ ${EVAL_AND_LOGGING_ARGS[@]} \
+ ${TRANSFORMER_ENGINE_ARGS[@]} \
+ ${PROFILE_ARGS[@]}
+ "
+
+USE_EPX=${USE_EPX:-0}
+
+# run cmd directly
+if [ $USE_EPX -eq 0 ]; then
+ echo $cmd
+ $cmd
+ exit $?
+fi
+
+# run cmd with fault tolerance
+source "${PATCH_HOME}/EXAMPLEs/deepseek-v2/deepseek-v2-lite/fault_tolerance_function.sh"
+ft_training "$cmd"
diff --git a/simu_tools/megatron_scripts/scripts/run_pretrain_llama3.sh b/simu_tools/megatron_scripts/scripts/run_pretrain_llama3.sh
new file mode 100644
index 0000000..24314ae
--- /dev/null
+++ b/simu_tools/megatron_scripts/scripts/run_pretrain_llama3.sh
@@ -0,0 +1,259 @@
+#!/bin/bash
+
+set -u
+WORK_HOME=${WORK_HOME:-"./"}
+PATCH_HOME=${PATCH_HOME:-"./"}
+MEGATRON_HOME=${MEGATRON_HOME:-"./"}
+EXAMPLE=${EXAMPLE:-"test"}
+HOSTFILE=${HOSTFILE:-"./hostfile"}
+TP_SIZE=${TP_SIZE:-1}
+PP_SIZE=${PP_SIZE:-1}
+MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
+GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-1}
+RDZV_ID=${RDZV_ID:-"0"}
+MODEL_TYPE=${MODEL_TYPE:-"llama3_8b"}
+TEST_TYPE=${TEST_TYPE:-"test"}
+NUM_LAYERS=${NUM_LAYERS:-1}
+DTYPE=${DTYPE:-"bf16"}
+OUTPUT_DIR=${OUTPUT_DIR:-"./output_${MODEL_TYPE}"}
+set +u
+
+if [ "${TEST_TYPE}" == "profile" ]; then
+ PROFILER_SAVE_PATH=${WORK_HOME}/profiler_result_${MODEL_TYPE}_${EXAMPLE}
+ rm -rf ${PROFILER_SAVE_PATH}
+ echo "Profiler save path: ${PROFILER_SAVE_PATH}"
+
+ TRAIN_ITERS=6
+ WARMUP_ITERS=1
+ PROFILE_ARGS=(
+ --profile
+ --profile-step-start 4
+ --profile-step-end 6
+ --use-pytorch-profiler
+ --profile-ranks 0 1 2 3 4 5 6 7
+ )
+else
+ PROFILE_ARGS=()
+ TRAIN_ITERS=10
+ WARMUP_ITERS=5
+ echo "Train iters: ${TRAIN_ITERS}, Warmup iters: ${WARMUP_ITERS}"
+fi
+export OMP_NUM_THREADS=4
+export CUDA_VISIBLE_DEVICES='0,1,2,3,4,5,6,7'
+export ACCELERATOR_BACKEND="cuda"
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+export PYTHONPATH=${MEGATRON_HOME}:${PATCH_HOME}:$PYTHONPATH
+# export NO_LOSS_REDUCE=1
+
+
+# CHECKPOINT_PATH=$WORK_HOME/checkpoints/$EXAMPLE
+# mkdir -p $CHECKPOINT_PATH
+
+# LOG_PATH=$OUTPUT_DIR/logs/$EXAMPLE
+# mkdir -p $LOG_PATH
+# cp $0 $LOG_PATH/
+TB_PATH=$OUTPUT_DIR/tboard/$EXAMPLE
+mkdir -p $TB_PATH
+# WB_PATH=$WORK_HOME/wandb/$EXAMPLE
+# mkdir -p $WB_PATH
+
+
+
+export NODE_ADDR=$(ip a|grep inet|grep -v 127.0.0.1|grep -v inet6|awk '{print $2;}'|tr -d "addr:"|head -n1 | cut -d '/' -f1)
+export GPUS_PER_NODE=8
+export NUM_NODES=$(cat $HOSTFILE | wc -l)
+export MASTER_ADDR=$(head -n1 $HOSTFILE | awk '{print $1;}')
+export NODE_RANK=$(awk '{ranks[$1]=(FNR-1);}END{print ranks["'$NODE_ADDR'"];}' $HOSTFILE)
+export MASTER_PORT=14388
+
+
+LOG_DIR=$OUTPUT_DIR/${DTYPE}_$EXAMPLE
+echo "Distributed log_dir: $LOG_DIR"
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --node_rank $NODE_RANK
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+ --log_dir $LOG_DIR
+ --redirects ${LOG_REDIRECTS_LEVEL:-3}
+)
+
+if [ "${MODEL_TYPE}" == "llama2_7b" ]; then
+ # llama2 7b
+ num_layer=$NUM_LAYERS
+ hidden_size=4096
+ ffn_hidden_size=11008
+ head_num=32
+ num_query_groups=32
+ vocab_size=32000
+elif [ "${MODEL_TYPE}" == "llama3_8b" ]; then
+ # llama3 8b
+ num_layer=$NUM_LAYERS
+ hidden_size=4096
+ ffn_hidden_size=14336
+ head_num=32
+ num_query_groups=8
+ vocab_size=128257
+elif [ "${MODEL_TYPE}" == "llama3_70b" ]; then
+ # llama3 70b
+ num_layer=$NUM_LAYERS # layer_num=60, 切层6
+ hidden_size=8192
+ ffn_hidden_size=28672
+ head_num=64
+ num_query_groups=8
+ vocab_size=128256
+elif [ "${MODEL_TYPE}" == "llama3_405b" ]; then
+ # llama3 405b
+ num_layer=$NUM_LAYERS # 405切1层
+ hidden_size=16384
+ ffn_hidden_size=53248
+ head_num=128
+ num_query_groups=16
+ vocab_size=128256
+else
+ echo "Model type ${MODEL_TYPE} not supported"
+ exit 1
+fi
+
+MODEL_ARGS=(
+ --num-layers ${num_layer}
+ --hidden-size ${hidden_size}
+ --ffn-hidden-size ${ffn_hidden_size}
+ --num-attention-heads ${head_num}
+ --group-query-attention
+ --num-query-groups ${num_query_groups}
+ --seq-length 4096
+ --max-position-embeddings 4096
+ --norm-epsilon 1e-5
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --disable-bias-linear
+ --position-embedding-type rope
+ --no-position-embedding
+ --swiglu
+ --normalization RMSNorm
+ --untie-embeddings-and-output-weights
+ --vocab-size ${vocab_size}
+ # --tp-comm-overlap-ag
+ # --tp_comm_overlap_rs
+ --disable-tp-comm-bulk-wgrad
+ --disable-tp-comm-bulk-dgrad
+ --disable-tp-comm-overlap-rs
+ --disable-tp-comm-overlap-ag
+ --disable-tp-comm-split-ag
+ --disable-tp-comm-split-rs
+)
+
+# 244140625 1T
+TRAINING_ARGS=(
+ --seed 42
+ --micro-batch-size $MICRO_BATCH_SIZE
+ --global-batch-size $GLOBAL_BATCH_SIZE
+ # --train-samples 24414062
+ --init-method-std 0.008
+ --use-mcore-models
+ --no-bias-dropout-fusion
+ # --no-bias-swiglu-fusion
+ --use-distributed-optimizer
+ --use-flash-attn
+ --sequence-parallel
+ --recompute-granularity full
+ --recompute-method block
+ --recompute-num-layers 0
+ --distributed-backend nccl
+ --mock-data
+)
+
+# --no-bias-swiglu-fusion
+# --no-rope-fusion
+# --no-gradient-accumulation-fusion
+# --transformer-impl local transformer_engine
+REGULARIZATION_ARGS=(
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.95
+ --clip-grad 1.0
+)
+
+WARMUP_STEPS=2000
+WARMUP_SAMPLES=$((WARMUP_STEPS * GLOBAL_BATCH_SIZE))
+
+LEARNING_RATE_ARGS=(
+ # --lr 1.5e-5
+ # --lr-decay-style cosine
+ # --lr-warmup-samples ${WARMUP_SAMPLES}
+ # --min-lr 1.5e-6
+ # --initial-loss-scale 65536
+ # --min-loss-scale 1.0
+ --initial-loss-scale 65536
+ --min-loss-scale 1.0
+
+ # 下面是ds 32b参数
+ --train-iters ${TRAIN_ITERS}
+ --lr-warmup-iters ${WARMUP_ITERS}
+ --lr 2.4e-4
+ --lr-decay-style cosine
+ --min-lr 1e-6
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size $TP_SIZE
+ --pipeline-model-parallel-size $PP_SIZE
+ # --decoder-last-pipeline-num-layers 14
+)
+
+MIXED_PRECISION_ARGS=(
+ --bf16
+ --attention-softmax-in-fp32
+ --no-masked-softmax-fusion
+ --accumulate-allreduce-grads-in-fp32
+)
+
+DATA_ARGS=(
+ --tokenizer-type NullTokenizer
+ --split 1
+)
+
+
+EVAL_AND_LOGGING_ARGS=(
+ --log-interval 1
+ --log-throughput
+ --save-interval 200000
+ --eval-interval 1000
+ # --save $CHECKPOINT_PATH
+ # --load $CHECKPOINT_PATH
+ --eval-iters 0
+ --tensorboard-dir $TB_PATH
+)
+
+if [ "$DTYPE" == "bf16" ]; then
+ TRANSFORMER_ENGINE_ARGS=(
+ --transformer-impl transformer_engine
+ )
+else
+ TRANSFORMER_ENGINE_ARGS=(
+ --transformer-impl transformer_engine
+ --fp8-format hybrid
+ --fp8-param-gather
+ )
+fi
+
+
+
+cmd="torchrun ${DISTRIBUTED_ARGS[@]} ./pretrain_gpt.py \
+ ${MODEL_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${REGULARIZATION_ARGS[@]}
+ ${LEARNING_RATE_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${MIXED_PRECISION_ARGS[@]}
+ ${DATA_ARGS[@]} \
+ ${EVAL_AND_LOGGING_ARGS[@]} \
+ ${TRANSFORMER_ENGINE_ARGS[@]} \
+ ${PROFILE_ARGS[@]}
+ "
+echo $cmd
+eval $cmd
diff --git a/simu_tools/megatron_scripts/stop_all.sh b/simu_tools/megatron_scripts/stop_all.sh
new file mode 100644
index 0000000..0b89a80
--- /dev/null
+++ b/simu_tools/megatron_scripts/stop_all.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+HOSTFILE=./hostfile
+NUM_NODES=$(grep -v '^#\|^$' $HOSTFILE | wc -l)
+echo "NUM_NODES: $NUM_NODES"
+
+hostlist=$(grep -v '^#\|^$' $HOSTFILE | awk '{print $1}' | xargs)
+for host in ${hostlist[@]}; do
+ ssh -f -n $host "pkill -f '/opt/conda/envs/py310/bin/torchrun'"
+ ssh -f -n $host "pkill -f '/usr/local/bin/torchrun'"
+ echo "$host is killed."
+done
diff --git a/simumax/core/base_struct.py b/simumax/core/base_struct.py
index fadbe8f..bead69a 100644
--- a/simumax/core/base_struct.py
+++ b/simumax/core/base_struct.py
@@ -5,6 +5,7 @@
from abc import ABC
from typing import List, Tuple, Dict
import time
+import types
import multiprocessing
try:
from mpi4py import MPI
@@ -15,11 +16,11 @@
import json
import os
from simumax.core.tensor import TensorSize
-from simumax.core.config import StrategyConfig, SystemConfig
+from simumax.core.config import StrategyConfig, SystemConfig, get_capture_graph_only, SIMU_DEBUG, TMP_PATH
from simumax.core.utils import get_point_name, path_convert_to_str, to_json_string
from simumax.core.utils import human_readable_bytes, human_readable_nums, human_readable_times
from simumax.core.utils import get_rank_group
-
+from simumax.core.graph import SimuONNXGraphBuilder
class RecomputeStatus:
NO_RECOMPUTE = "no_recompute"
@@ -293,36 +294,41 @@ class ModuleMemoryInfo:
2. Gradient memory usage
3. State memory usage
"""
-
- weight_bytes: int = 0
- grad_bytes: int = 0
- state_bytes: int = 0
+ weight_numel: int = 0
+ dense_weight_bytes: int = 0
+ dense_grad_bytes: int = 0
+ dense_state_bytes: int = 0
+ moe_weight_numel: int = 0
moe_weight_bytes: int = 0
moe_grad_bytes: int = 0
moe_state_bytes: int = 0
@property
def all(self):
- return self.weight_bytes + self.grad_bytes + self.state_bytes + self.moe_weight_bytes + self.moe_grad_bytes + self.moe_state_bytes
+ return self.dense_weight_bytes + self.dense_grad_bytes + self.dense_state_bytes + self.moe_weight_bytes + self.moe_grad_bytes + self.moe_state_bytes
@property
def all_state_bytes(self):
- return self.state_bytes + self.moe_state_bytes
+ return self.dense_state_bytes + self.moe_state_bytes
@property
def all_weight_bytes(self):
- return self.weight_bytes + self.moe_weight_bytes
+ return self.dense_weight_bytes + self.moe_weight_bytes
+
+ @property
+ def all_weight_numel(self):
+ return self.weight_numel + self.moe_weight_numel
@property
def all_grad_bytes(self):
- return self.grad_bytes + self.moe_grad_bytes
+ return self.dense_grad_bytes + self.moe_grad_bytes
def _format_repr_info(self):
attributes = {
"all": self.all,
- "weight_bytes": self.weight_bytes,
- "grad_bytes": self.grad_bytes,
- "state_bytes": self.state_bytes,
+ "weight_bytes": self.dense_weight_bytes,
+ "grad_bytes": self.dense_grad_bytes,
+ "state_bytes": self.dense_state_bytes,
"moe_weight_bytes": self.moe_weight_bytes,
"moe_grad_bytes": self.moe_grad_bytes,
"moe_state_bytes": self.moe_state_bytes
@@ -345,9 +351,11 @@ def __add__(self, other):
f"Unsupported operand type for +: ModuleMemoryInfo and {type(other)}"
)
return ModuleMemoryInfo(
- weight_bytes=self.weight_bytes + other.weight_bytes,
- grad_bytes=self.grad_bytes + other.grad_bytes,
- state_bytes=self.state_bytes + other.state_bytes,
+ weight_numel=self.weight_numel + other.weight_numel,
+ dense_weight_bytes=self.dense_weight_bytes + other.dense_weight_bytes,
+ dense_grad_bytes=self.dense_grad_bytes + other.dense_grad_bytes,
+ dense_state_bytes=self.dense_state_bytes + other.dense_state_bytes,
+ moe_weight_numel = self.moe_weight_numel + other.moe_weight_numel,
moe_weight_bytes=self.moe_weight_bytes + other.moe_weight_bytes,
moe_grad_bytes=self.moe_grad_bytes + other.moe_grad_bytes,
moe_state_bytes=self.moe_state_bytes + other.moe_state_bytes
@@ -379,6 +387,12 @@ class ModuleCostInfo:
def fwd_time(self):
return self.fwd_compute_time + self.fwd_net_exposed_time
+ # def all_time(self):
+ # return self.fwd_time + self.fwd_net_time + self.recompute_compute_time + self.recompute_net_time + self.bwd_time + self.bwd_net_time
+ @property
+ def all_time(self):
+ return self.fwd_time + self.fwd_net_time + self.bwd_time + self.bwd_net_time
+
@property
def recompute_time(self):
return self.recompute_compute_time + self.recompute_net_exposed_time
@@ -394,6 +408,10 @@ def bwd_time(self):
@property
def bwd_net_time(self):
return self.bwd_grad_w_net_time + self.bwd_grad_act_net_time
+
+ @property
+ def net_time(self):
+ return self.fwd_net_time + self.bwd_net_time + self.recompute_net_time
def get_all_costs(self):
"""Get all cost of the model in forward and backward(bwd_act, bwd_w) pass"""
@@ -556,11 +574,12 @@ class MetaModule(BaseModel, metaclass = PostInitMeta):
"""
dtype_to_element_size = {"fp32": 4, "fp16": 2, "bf16": 2, "fp8": 1}
-
+ id_counter = 0
def __init__(self, strategy:StrategyConfig, system:SystemConfig, specific_name='', parent_module = None) -> None:
super().__init__(specific_name)
self.strategy = strategy
self.system = system
+ self.offload_inputs = False
self.children_ordered_module:List[MetaModule] = []
self.children_modules:List[MetaModule] = [] # children modules是所有子模块的列表（无序）
@@ -568,7 +587,8 @@ def __init__(self, strategy:StrategyConfig, system:SystemConfig, specific_name='
self.default_dtype = strategy.dtype
self._init_strategy = False
self.input_info = None
- self.cache_info = []
+ self.output_info_ = None
+ # self.cache_info = []
self.enable_recompute = False
self.recompute_granularity = "full"
self.parent_module:MetaModule = parent_module
@@ -585,11 +605,15 @@ def __init__(self, strategy:StrategyConfig, system:SystemConfig, specific_name='
self.is_recompute_forward_finished = False
self.full_name = "self"
self.name = ''
+ self.call_idx = -1
# for Selective recompute strategy
self.all_recompute_nodes:List[MetaModule] = []
self.all_leaf_nodes:List[MetaModule] = []
self.status_ready = False
+ self.is_variance_node = False
+ self.id = MetaModule.id_counter
+ MetaModule.id_counter += 1
def __post_init__(self):
self.is_leaf_module = self.set_children_modules()
@@ -607,6 +631,15 @@ def set_children_modules(self):
self.children_modules_names[member] = name
return is_leaf
+ def set_variance_node(self, is_variance_node:bool):
+ if self.strategy.recompute_variance:
+ self.is_variance_node = is_variance_node
+ @property
+ def output_info(self):
+ if self.output_info_ is None:
+ self.output_info_ = self.create_output_info()
+ return self.output_info_
+
def set_leaf_full_name(self, parent_name:str):
for child, name in self.children_modules_names.items():
child.full_name = parent_name + '.' + name
@@ -733,18 +766,6 @@ def compute_end2end_time(self, compute_time, mem_time):
def all_input_element_num(self):
res = 0
- # for x in self.input_info.tensors:
- # if hasattr(self, "a_element_size"):
- # element_size = self.a_element_size
- # else:
- # element_size = self.element_size
- # res += x.numel() * element_size
- # return res
-
- # if hasattr(self, "a_element_size"):
- # element_size = self.a_element_size # for fp8 gemm
- # else:
- # element_size = self.element_size
if isinstance(self.input_info, InputOutputInfo):
input_info = [self.input_info]
@@ -757,33 +778,7 @@ def all_input_element_num(self):
elif isinstance(ii, TensorSize):
res += ii.get_memory_size()
return res
- def all_output_shape(self):
- res = 0
- # for x in self.input_info.tensors:
- # if hasattr(self, "a_element_size"):
- # element_size = self.a_element_size
- # else:
- # element_size = self.element_size
- # res += x.numel() * element_size
- # return res
-
- # if hasattr(self, "a_element_size"):
- # element_size = self.a_element_size # for fp8 gemm
- # else:
- # element_size = self.element_size
- shapes = []
- # element_size = self.element_size
- if isinstance(self.output_info, InputOutputInfo):
- output_info = [self.output_info]
- else:
- output_info = self.output_info
- for oi in output_info:
- if isinstance(oi, InputOutputInfo):
- for x in oi.tensors:
- shapes.append(x.shape)
- elif isinstance(oi, TensorSize):
- shapes.append(oi.shape)
- return shapes
+
def all_output_element_num(self):
res = 0
# element_size = self.element_size
@@ -802,17 +797,11 @@ def all_output_element_num(self):
def set_input_state_info(self, input_info: InputOutputInfo):
# self.input_info = deepcopy(input_info)
self.input_info = input_info # reference assignments are allowed here
-
- def save_for_backward(self, input_info: InputOutputInfo):
- self.cache_info.append(input_info) # reference assignments are allowed here
def set_path_debug_context(self, path_debug_context: PathDebugContext):
self.path_debug_context = deepcopy(path_debug_context)
- @property
- def output_info(self):
-
- # raise NotImplementedError(f"{self.__class__.__name__} 没有实现 output_info 属性")
+ def create_output_info(self):
return InputOutputInfo([])
# =========================
@@ -842,12 +831,20 @@ def _comp_act_info(self):
# leaf module act info is the same with recompute,
# because _act_info_with_recomp is used to distinguish
# the case of recompute in the combined module
+ if self.is_variance_node:
+ # print("Warning: variance node change peak")
+ # self._act_info.activation_mem_cache = 0
+ # self._act_info.bwd_peak_mem_no_cache += self._act_info.activation_mem_cache
+ pass
self._act_info_with_recomp = deepcopy(self._act_info)
+ else:
+ for module in self.children_ordered_module:
+ self._act_info.activation_mem_cache = self._act_info.activation_mem_cache + module._act_info.activation_mem_cache
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_model_info(self):
if len(self.children_ordered_module) > 0:
@@ -856,22 +853,6 @@ def _comp_model_info(self):
else:
self._comp_leaf_model_info_impl()
-
- def _comp_input_cache_size(self):
- def get_cache_size(cache):
- if cache is None:
- return 0
- if isinstance(cache, InputOutputInfo):
- # return sum([t.try_get_allocated_memory() for t in cache.tensors])
- return sum([t.get_memory_size() for t in cache.tensors])
- elif isinstance(cache, TensorSize):
- return cache.get_memory_size()
- else:
- raise TypeError(f"{self.__class__.__name__} cache_info should be InputOutputInfo or TensorSize, but got {type(cache)}")
-
- all_cache_size = [get_cache_size(cache) for cache in self.cache_info]
- return sum(all_cache_size)
-
# =========================
# Compute or Communicate Related
# =========================
@@ -899,6 +880,13 @@ def _comp_compute_info(self):
self._comp_leaf_flops_info()
self._comp_leaf_mem_accessed_info()
self._comp_leaf_intra_net_info()
+ if self.strategy.recompute_variance and self.is_variance_node:
+ self._compute_info.recompute_accessed_mem = 0
+ self._compute_info.recompute_flops = 0
+ self._cost_info.recompute_net_time = 0
+ self._cost_info.recompute_net_exposed_time = 0
+ if SIMU_DEBUG:
+ print(f"- {self.full_name} is variance node, recompute_accessed_mem and recompute_flops are set to 0")
def _comp_cost_info(self):
if len(self.children_ordered_module) > 0:
@@ -911,7 +899,8 @@ def _comp_cost_info(self):
bwd_grad_act_op="default",
bwd_grad_w_op="default",
enable_recompute=self.enable_recompute,
- )
+ )
+
if (
self.path_debug_context
and self.path_debug_context.target_point is not None
@@ -952,40 +941,6 @@ def set_details(
"io_details" : deepcopy(io_details),
}
- def get_input_shapes_desc2(self, stage):
- m, n, k = 0, 0, 0
- weight = self.get_weight()
- if weight:
- weight_shape = [weight.shape] # [out, in]
- weight_T_shape = [weight.T.shape] # [in, out]
- else:
- weight_shape = []
- weight_T_shape = []
-
- if isinstance(self, LinearBase):
- if stage == 'fwd':
- input_shapes = self.input_info.shapes + weight_T_shape
- lhs = self.input_info.tensors[0]
- rhs = weight.T
- elif stage == 'bwd_grad_act':
- output_shape = self.output_info.shapes if isinstance(self.output_info, InputOutputInfo) else [self.output_info.shape]
- input_shapes = output_shape + weight_T_shape
- elif stage == 'bwd_grad_w':
- d_w_lhs_shape = [self.input_info.tensors[0].transpose(-1,-2).shape]
- input_shapes = d_w_lhs_shape + (self.output_info.shapes if isinstance(self.output_info, InputOutputInfo) else [self.output_info.shape])
- else:
- raise ValueError(f"Unknown stage: {stage}")
-
- input_shapes = [[int(i) for i in shape] for shape in input_shapes]
- lhs_shape = list(input_shapes[0])
- rhs_shape = list(input_shapes[1])
- lhs_shape = [int(x) for x in lhs_shape]
- rhs_shape = [int(x) for x in rhs_shape]
- shapes_desc = str(lhs_shape) + ' x ' + str(rhs_shape)
- else:
- shapes_desc = ""
- return shapes_desc
-
def get_input_shapes_desc(self, stage):
if isinstance(self, LinearBase):
bmnk_info = self.get_gemm_bmnk(stage)
@@ -1005,7 +960,7 @@ def _comp_cost_info_impl(
):
def compute_details(op_name, stage, flops, accessed_mem):
#compute_details include compute time, tflops of accelerator, flops of current op, etc.
- compute_details = self.system.compute_op_accuracy_time2(op_name, flops, shape_desc= self.get_input_shapes_desc(stage), reture_detail=True)
+ compute_details = self.system.compute_op_accuracy_time(op_name, flops, shape_desc= self.get_input_shapes_desc(stage), reture_detail=True)
# io_details include io time, gbps of accelerator, io size of current op, etc.
io_details = self.system.compute_mem_access_time(op_name,
@@ -1027,40 +982,45 @@ def compute_details(op_name, stage, flops, accessed_mem):
self._cost_info.bwd_grad_act_time = compute_details(bwd_grad_act_op, 'bwd_grad_act', self._compute_info.bwd_grad_act_flops, self._compute_info.bwd_grad_act_accessed_mem)
self._cost_info.bwd_grad_w_time = compute_details(bwd_grad_w_op, 'bwd_grad_w', self._compute_info.bwd_grad_w_flops, self._compute_info.bwd_grad_w_accessed_mem)
- if enable_recompute:
- self._cost_info.recompute_compute_time = self._cost_info.fwd_time
-
- if (
- self.path_debug_context
- and self.path_debug_context.target_point is not None
- ):
- # get the parent path of the current module
- path = get_point_name(
- parent=self.parent, current=self.current, sep=" -> "
- )
- if path in self.path_debug_context.target_point:
- file_path = f'{TMP_PATH}/cost_log.json'
- os.makedirs(TMP_PATH, exist_ok=True)
- if os.path.exists(file_path):
- with open(file_path, 'r', encoding='utf-8') as file:
- try:
- existing_data = json.load(file)
- except json.JSONDecodeError:
- existing_data = {}
- else:
- existing_data = {}
- existing_data.update(
- {path:{"cost_F": self._cost_info.fwd_compute_time,
- "cost_B": self._cost_info.bwd_grad_act_time,
- "cost_W": self._cost_info.bwd_grad_w_time,
- "recompute_F": self._cost_info.recompute_compute_time,
- "net_F": self._cost_info.fwd_net_time,
- "net_B": self._cost_info.bwd_net_time,
- }
- }
- )
- with open(file_path, 'w', encoding='utf-8') as file:
- json.dump(existing_data, file, indent=4, ensure_ascii=False)
+ self._cost_info.recompute_compute_time = self._cost_info.fwd_time if self.enable_recompute else 0
+
+ if self.enable_recompute and self.is_variance_node:
+ self._cost_info.recompute_compute_time = 0
+ if SIMU_DEBUG:
+ # if 1:
+ print(f'%% {self.name} is variance node, recompute_compute_time is 0')
+
+ # if (
+ # self.path_debug_context
+ # and self.path_debug_context.target_point is not None
+ # ):
+ # # get the parent path of the current module
+ # path = get_point_name(
+ # parent=self.parent, current=self.current, sep=" -> "
+ # )
+ # if path in self.path_debug_context.target_point:
+ # file_path = f'{TMP_PATH}/cost_log.json'
+ # os.makedirs(TMP_PATH, exist_ok=True)
+ # if os.path.exists(file_path):
+ # with open(file_path, 'r', encoding='utf-8') as file:
+ # try:
+ # existing_data = json.load(file)
+ # except json.JSONDecodeError:
+ # existing_data = {}
+ # else:
+ # existing_data = {}
+ # existing_data.update(
+ # {path:{"cost_F": self._cost_info.fwd_compute_time,
+ # "cost_B": self._cost_info.bwd_grad_act_time,
+ # "cost_W": self._cost_info.bwd_grad_w_time,
+ # "recompute_F": self._cost_info.recompute_compute_time,
+ # "net_F": self._cost_info.fwd_net_time,
+ # "net_B": self._cost_info.bwd_net_time,
+ # }
+ # }
+ # )
+ # with open(file_path, 'w', encoding='utf-8') as file:
+ # json.dump(existing_data, file, indent=4, ensure_ascii=False)
# =========================
# Agg Related
@@ -1102,6 +1062,7 @@ def forward(self, input_info: InputOutputInfo, path_debug_context: PathDebugCont
def __call__(
self, input_info: InputOutputInfo, path_debug_context: PathDebugContext
) -> InputOutputInfo:
+ is_capture_only = get_capture_graph_only()
if isinstance(input_info, TensorSize):
input_info = InputOutputInfo([input_info])
@@ -1129,34 +1090,30 @@ def __call__(
self.current = current_repr
self.current_full_module_path = get_point_name(parent=self.parent, current=self.current, sep=" -> ") #FIXME(sherry): path_debug_context is deepcopy to module. How to modify the parent of the temporary variable and pass it to the next module?
-
-
# call once, return all fwd, bwd info
self._pre_op()
output_info = None
-
- self.call_forward_pre_hook(input_info)
-
- # Save the input info for backward
- if self.cache_inputs:
- self.save_for_backward(input_info)
if not self.is_leaf_module:
output_info = self.forward(input_info, self.path_debug_context)
+ else:
+ output_info = output_info if output_info else self.output_info # output_info = None, return leaf output
+ if is_capture_only:
+ graph_builder = SimuONNXGraphBuilder()
+ graph_builder.add_node(op = self,
+ op_type = self.__class__.__name__,
+ inputs = input_info.tensors if isinstance(input_info, InputOutputInfo) else [input_info],
+ outputs = output_info.tensors if isinstance(output_info, InputOutputInfo) else [output_info]
+ )
- output_info = output_info if output_info else self.output_info # output_info = None, return leaf output
+ if not is_capture_only:
+ # aggregate the info or compute the leaf info
+ self._comp_model_info() #static model memory usage
+ self._comp_act_info() #activation
+ self._comp_compute_info()
+ self._post_op()
+ self._comp_cost_info()
- # Save the output info for backward
- if self.cache_outputs:
- self.save_for_backward(output_info)
- self.call_forward_post_hook(input_info, output_info)
-
- # aggregate the info or compute the leaf info
- self._comp_model_info() #static model memory usage
- self._comp_act_info() #activation
- self._comp_compute_info()
- self._post_op()
- self._comp_cost_info()
self._info_ready = True
if isinstance(output_info, InputOutputInfo) and len(output_info.tensors) == 1:
@@ -1231,22 +1188,33 @@ def _addindent(s_, numSpaces):
main_str += extra_lines[0]
else:
main_str += "\n " + "\n ".join(lines) + "\n"
- main_str += ")"
+
module = self
# TODO(sherry): delete this, for debug
- debug = False
- if self.is_leaf_module and debug:
- main_str += f"\n\t\t-- fwd_call_idx={module._act_info.fwd_idx}, all_input={module.all_input_element_num()/1024/1024:.2f} MB, all_ouput={module.all_output_element_num()/1024/1024:.2f} MB, output_shape={module.all_output_shape()}, activation_mem_cache={module._act_info_with_recomp.activation_mem_cache/1024/1024:.2f} MB, "
- main_str += f"\n\t\t-- recompute_status={module.recompute_status}"
- cost_info = module._cost_info
- main_str += f"\t\t{self.name} fwd_cost={cost_info.fwd_time+cost_info.fwd_net_time:.2f} ms-[compute={cost_info.fwd_compute_time*1000:.2f} us, net={cost_info.fwd_net_time*1000:.2f} us], bwd_cost={cost_info.bwd_time+cost_info.bwd_net_time:.2f} ms-[compute={cost_info.bwd_compute_time*1000:.2f} us, net={cost_info.bwd_net_time*1000:.2f} us]"
- # main_str += f"\t\t{cost_info}"
- # , fwd_flops={module._compute_info.fwd_flops:.0f}, bwd_flops={module._compute_info.bwd_flops:.0f}
- # fwd_activation={module._act_info.activation_mem_cache/1024/1024:.2f} MB"
+
+ main_str += ")"
+
+ show_details = True
+ if show_details:
+ cost_info = module._cost_info
+ main_str += f"\n\t1. cost: (total_time={cost_info.all_time:.2f} ms, fwd_details=(sum={cost_info.fwd_time+cost_info.fwd_net_time:.2f} ms, compute={cost_info.fwd_compute_time*1000:.2f} us, net={cost_info.fwd_net_time*1000:.2f} us), bwd_details=(sum={cost_info.bwd_time+cost_info.bwd_net_time:.2f} ms, compute={cost_info.bwd_compute_time*1000:.2f} us, net={cost_info.bwd_net_time*1000:.2f} us), variance_node={self.is_variance_node} flops={sum(module._compute_info.get_all_flops())/1e12:.2f} T) "
+
+ module_info = module._model_info
+ main_str += f"\n\t2. memory: (d_w={module_info.dense_weight_bytes}, d_g={module_info.dense_grad_bytes}, m_w={module_info.moe_weight_bytes}, m_g={module_info.moe_grad_bytes})"
+
return main_str
+class RecomputeBreakModule(MetaModule):
+ def __init__(self, strategy, system, specific_name='', parent_module=None):
+ super().__init__(strategy, system, specific_name, parent_module=parent_module)
+ self.enable_recompute = False
+
+ # TODO(sherry): no memory and not cost. Need to be implemented
+ def create_output_info(self):
+ output_info = InputOutputInfo(tensors=[t.new() for t in self.input_info.tensors])
+ return output_info
class LinearBase(MetaModule):
def __init__(self, input_size, output_size, strategy, system, specific_name='', parent_module=None):
super().__init__(strategy, system, specific_name, parent_module)
@@ -1333,8 +1301,13 @@ def __init__(self, local_expert_num, input_size: int, output_size: int, strateg
self.local_expert_num = local_expert_num
def get_input_shapes_desc(self, stage):
+ assert self.input_info.tensors[0].size(0) % self.local_expert_num == 0, f'input size {self.input_info.tensors[0].size(0)} is not divisible by local_expert_num {self.local_expert_num} {self.strategy.parallelism}'
num_tokens = self.input_info.tensors[0].size(0) // self.local_expert_num
shape_str = f'ng={self.local_expert_num}, M={num_tokens}, N={self.output_size}, K={self.input_size}'
+
+ dtype_str = f", dtype={'fp8' if self.strategy.fp8 else 'bf16'}, out_dtype=bf16, main_grad_dtype={'bf16' if self.strategy.grad_reduce_in_bf16 else 'fp32'}"
+ # if self.strategy.fp8:
+ shape_str += dtype_str
if stage == 'fwd':
shape_str += ', stage=fwd, grad=False, accumulate=False, use_split_accumulator=False, single_output=True'
elif stage == 'bwd_grad_act':
diff --git a/simumax/core/config.py b/simumax/core/config.py
index cd498ca..d6ffe0b 100644
--- a/simumax/core/config.py
+++ b/simumax/core/config.py
@@ -1,15 +1,19 @@
"""Configuration classes for SimuMax """
import os
import time
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass, asdict, field
from typing import Optional, Dict, Any
+from collections import OrderedDict
import json
import copy
import math
import warnings
+import re
from simumax.core.utils import to_json_string
+capture_graph_only = False
+ENABLE_SIMU_GRAPH = int(os.environ.get("ENABLE_SIMU_GRAPH", "0"))
SIMU_CHECK = int(os.environ.get("SIMU_CHECK", "0"))
SIMU_DEBUG = int(os.environ.get('SIMU_DEBUG', '0'))
if SIMU_CHECK:
@@ -26,6 +30,47 @@
)
+def set_capture_graph_only(value: bool):
+ global capture_graph_only
+ capture_graph_only = value
+
+def get_capture_graph_only():
+ return capture_graph_only
+
+class ParameterExtractor:
+ def __init__(self, param_patterns: Dict[str, Any]):
+ # 定义参数模式及其默认值
+ self.param_patterns = param_patterns
+
+ def extract_parameters(self, input_string):
+ """从字符串中提取所有参数，缺失的使用默认值"""
+ parameters = {}
+
+ for param_name, (pattern, default_value) in self.param_patterns.items():
+ match = re.search(pattern, input_string)
+ if match:
+ parameters[param_name] = int(match.group(1))
+ elif default_value is not None:
+ parameters[param_name] = default_value
+ print(f"警告: 未找到参数 {param_name}, 使用默认值 {default_value}")
+
+ return parameters
+
+ def extract_single_parameter(self, input_string, param_name, default_value=None):
+ """提取单个参数"""
+ if param_name not in self.param_patterns:
+ raise ValueError(f"未知参数: {param_name}")
+
+ pattern, default = self.param_patterns[param_name]
+ if default_value is not None:
+ default = default_value
+
+ match = re.search(pattern, input_string)
+ if match:
+ return int(match.group(1))
+ else:
+ print(f"警告: 未找到参数 {param_name}, 使用默认值 {default}")
+ return default
@dataclass
class Config:
"""
@@ -84,27 +129,42 @@ def init_from_config_file(cls, config_file: str):
@dataclass
class AttentionRecomputeConfig(Config):
- input_norm_recompute:bool = False
- qkv_norm_recompute:bool = False
- qkv_recompute:bool = False
- attn_recompute:bool = False
+ # input_norm_recompute:bool = False
+ # qkv_norm_recompute:bool = False
+ # qkv_recompute:bool = False
+ # attn_recompute:bool = False
+ # out_recompute:bool = False
+
+ input_layernorm_recompute:bool = False
+
+ q_down_recompute:bool = False
+ kv_down_recompute:bool = False
+ q_up_recompute:bool = False
+ kv_up_recompute:bool = False
+
+ q_layernorm_recompute:bool = False
+ kv_layernorm_recompute:bool = False
+
+ rope_recompute:bool = False
+ core_attn_recompute:bool = False
+
out_recompute:bool = False
+ def set_all_status(self, status:bool):
+ self.__dict__.update({k:status for k in self.__dict__.keys()})
+
@property
def is_recompute_all(self):
- return (self.input_norm_recompute and
- self.qkv_norm_recompute and
- self.qkv_recompute and
- self.attn_recompute and
- self.out_recompute)
+ return all(self.__dict__.values())
@dataclass
class MLPRecomputeConfig(Config):
pre_mlp_norm_recompute:bool = False
- linear_recompute:bool = False
+ shared_linear_recompute:bool = False
+ linear_recompute:bool = False # Noraml MLP and grouped MLP
router_recompute:bool = False
permutation_recompute:bool = False
-
+
@property
def is_recompute_all(self):
return (self.pre_mlp_norm_recompute and
@@ -120,11 +180,12 @@ class StrategyConfig(Config):
seq_len: Optional[int] = None
micro_batch_size: Optional[int] = None
micro_batch_num: Optional[int] = None
- dtype: Optional[int] = None
+ dtype: Optional[int] = 'bf16'
fp8: Optional[bool] = False
+
# dist strategy
- world_size: Optional[int] = None
- tp_size: int = None
+ world_size: Optional[int] = 8
+ tp_size: int = 1
pp_size: int = 1
ep_size: int = 1
etp_size: int = 1
@@ -133,7 +194,11 @@ class StrategyConfig(Config):
num_layers_in_last_pipeline_stage: Optional[int] = None
account_for_embedding_in_pipeline_split: bool = False
account_for_loss_in_pipeline_split: bool = False
+
+ # memory optimization
grad_reduce_in_bf16: bool = False
+ cache_groupgemm_col_fp8_inputs: Optional[bool] = False
+ offload_groupgemm_col_inputs: Optional[bool] = False
attn_recompute: bool = False
mla_rms_recompute: bool = False
@@ -142,9 +207,8 @@ class StrategyConfig(Config):
enable_sequence_parallel: bool = True
interleaving_size: int = 1
- zero_state: int = 0
+ zero_state: int = 1
- no_sync: bool = True
attention_sparse_ratio: float = (
0.0 # 0.0 means dense attention; 0.5 means compute optimize for causal attention
)
@@ -153,10 +217,10 @@ class StrategyConfig(Config):
use_accm_weight:bool = True # TODO(sherry): if True, No need to generate temporary variables of weight
# recompute
- enable_recompute: bool = False
- skip_ckpt_micro_batch_num: int = 1
+ enable_recompute: bool = True
recompute_granularity: Optional[str] = None
recompute_layer_num: int = 0
+ recompute_variance: bool = False
# fused kernel
use_flash_sdp: bool = True
@@ -172,10 +236,12 @@ class StrategyConfig(Config):
dp_net: Optional[str] = "auto"
ep_net: Optional[str] = "auto"
etp_net: Optional[str] = "auto"
+ edp_net: Optional[str] = "auto"
- mem_factor: float = 0.94
+ # Megatron related
+ dispatch_probs: bool = False # The new version of Megatron combines probs in Silu after Groupgemm1 in ExpertMLP
- recompute_breakpoints = ["LLMBlock", "MLAAttention", "MLP"]
+ mem_factor: float = 0.94
valid_recompute_granularity = [
"full_block",
@@ -184,6 +250,37 @@ class StrategyConfig(Config):
"sdp_only",
"selective_recompute"
]
+
+ @classmethod
+ def init_from_format_strings(cls, strs):
+ """
+ Docstring for init_from_format_strings
+ parse format like:
+ 查找
+ seq{self.seq_len}.mbs{self.micro_batch_size}.mbc{self.micro_batch_num}.gbs{self.global_batch_size} tp{self.tp_size}.ep{self.ep_size}.pp{self.pp_size}.dp{self.dp_size}.etp{self.etp_size}.edp{self.edp_size}, world_size:{self.world_size}
+
+ :param cls: Description
+ :param strs: Description
+ :return: Description
+ :rtype: Any
+ """
+ param_patterns = {
+ 'seq_len': (r'seq(\d+)', 4096),
+ 'micro_batch_size': (r'mbs(\d+)', 1),
+ 'micro_batch_num': (r'mbc(\d+)', 1),
+ 'global_batch_size': (r'gbs(\d+)', 1),
+ 'tp_size': (r'tp(\d+)', 1),
+ 'ep_size': (r'ep(\d+)', 1),
+ 'pp_size': (r'pp(\d+)', 1),
+ 'world_size': (r'world_size:(\d+)', 8)
+ }
+ extractor = ParameterExtractor(param_patterns=param_patterns)
+ params = extractor.extract_parameters(strs)
+ global_batch_size = params.pop('global_batch_size')
+ strategty = StrategyConfig(**params)
+ strategty.reset_global_batch_size(global_batch_size)
+ return strategty
+
@property
def shard_size(self):
return self.pp_size * self.tp_size
@@ -206,78 +303,78 @@ def edp_size(self):
def parallelism(self):
return f'seq{self.seq_len}.mbs{self.micro_batch_size}.mbc{self.micro_batch_num}.gbs{self.global_batch_size} tp{self.tp_size}.ep{self.ep_size}.pp{self.pp_size}.dp{self.dp_size}.etp{self.etp_size}.edp{self.edp_size}, world_size:{self.world_size}'
+ @property
+ def is_recompute(self):
+ is_full_recompute = self.recompute_layer_num > 0 and self.recompute_granularity == 'full_block'
+ is_selective_recompute = self.recompute_layer_num > 0 and self.recompute_granularity == 'selective_recompute' and any([self.attn_recompute, self.mla_rms_recompute, self.mlp_recompute, self.mlp_rms_recompute])
+ return self.enable_recompute and (is_full_recompute or is_selective_recompute)
+
@property
def recompute_status(self):
is_full_recompute = self.recompute_layer_num > 0 and self.recompute_granularity == 'full_block'
is_selective_recompute = self.recompute_layer_num > 0 and self.recompute_granularity == 'selective_recompute' and any([self.attn_recompute, self.mla_rms_recompute, self.mlp_recompute, self.mlp_rms_recompute])
- if not is_full_recompute and not is_selective_recompute:
+ if not self.is_recompute:
return 'No Recompute'
if is_full_recompute:
return f"{self.recompute_granularity}, recompute_layer_num={self.recompute_layer_num}"
+ elif is_selective_recompute:
+ return f'{self.recompute_granularity}, recompute_layer_num={self.recompute_layer_num}, attn={self.attn_recompute}, attn_rms={self.mla_rms_recompute}, mlp={self.mlp_recompute}, mlp_rms={self.mlp_rms_recompute}, recompute_variance={self.recompute_variance}'
else:
- return f'{self.recompute_granularity}, recompute_layer_num={self.recompute_layer_num}, attn={self.attn_recompute}, attn_rms={self.mla_rms_recompute}, mlp={self.mlp_recompute}, mlp_rms={self.mlp_rms_recompute}'
+ return 'Unknown Recompute Status'
@property
def net(self):
return f"pp_net={self.pp_net}, tp_net={self.tp_net}, dp_net={self.dp_net}, ep_net={self.ep_net}, etp_net={self.etp_net}"
def parse_attention_recompute(self, layer_idx):
- if self.recompute_granularity is None:
+ if self.recompute_granularity is None or layer_idx >= self.recompute_layer_num:
return AttentionRecomputeConfig()
-
+ conf = AttentionRecomputeConfig()
if self.recompute_granularity == "full_block":
- input_norm_recompute = True
- qkv_norm_recompute = True
- qkv_recompute = True
- attn_recompute = True
- out_recompute = True
+ conf.set_all_status(True)
elif self.recompute_granularity == "attn_only":
- input_norm_recompute = False
- qkv_norm_recompute = True # TODO(sherry): check this, theoretically, it should be enabled, but the old version does not
- qkv_recompute = True # attn only
- attn_recompute = True # attn only
- out_recompute = True # attn only
+ conf.q_down_recompute = True
+ conf.kv_down_recompute = True
+ conf.q_up_recompute = True
+ conf.kv_up_recompute = True
+ conf.q_layernorm_recompute = True
+ conf.kv_layernorm_recompute = True
+ conf.rope_recompute = True
+ conf.core_attn_recompute = True
+ conf.out_recompute = True
elif self.recompute_granularity == "sdp_only":
- input_norm_recompute = False
- qkv_norm_recompute = False
- qkv_recompute = False
- attn_recompute = True # sdp only
- out_recompute = False
+ raise NotImplementedError("sdp_only is not implemented yet")
+
elif self.recompute_granularity == "selective_recompute":
- # selective_recompute support:
- # attn_recompute: bool = False
- # mla_rms_recompute: bool = False
- # mlp_recompute: bool = False
- # mlp_rms_recompute: bool = False
- input_norm_recompute = self.mla_rms_recompute # normalization before attention
-
if self.mla_rms_recompute:
assert self.attn_recompute, "mla_rms_recompute requires attn_recompute"
- qkv_norm_recompute = self.mla_rms_recompute or self.attn_recompute # qkv norm recompute made with attn recompute
- qkv_recompute = self.mla_rms_recompute or self.attn_recompute
- attn_recompute = self.attn_recompute
- out_recompute = False
+ conf.input_layernorm_recompute = self.mla_rms_recompute
+ conf.q_down_recompute = self.mla_rms_recompute
+ conf.kv_down_recompute = self.mla_rms_recompute
+ conf.q_up_recompute = self.attn_recompute
+ conf.kv_up_recompute = self.attn_recompute
+ conf.q_layernorm_recompute = self.attn_recompute
+ conf.kv_layernorm_recompute = self.attn_recompute
+ conf.rope_recompute = self.attn_recompute
+ conf.core_attn_recompute = self.attn_recompute
+ conf.out_recompute = False
else:
raise ValueError("Invalid recompute_granularity")
- enable_layer_recompute = (
- layer_idx < self.recompute_layer_num
- )
- return AttentionRecomputeConfig(input_norm_recompute and enable_layer_recompute,
- qkv_norm_recompute and enable_layer_recompute,
- qkv_recompute and enable_layer_recompute,
- attn_recompute and enable_layer_recompute,
- out_recompute and enable_layer_recompute)
+
+ return conf
def parse_mlp_recompute(self, layer_idx):
- if self.recompute_granularity is None:
+ if self.recompute_granularity is None or layer_idx >= self.recompute_layer_num:
return MLPRecomputeConfig()
if self.recompute_granularity == "full_block":
pre_mlp_norm_recompute = True
linear_recompute = True
+ shared_linear_recompute = True
router_recompute = True
permutation_recompute = True
elif self.recompute_granularity in ["attn_only", "sdp_only"]:
pre_mlp_norm_recompute = False
+ shared_linear_recompute = False
linear_recompute = False
router_recompute = False
permutation_recompute = False
@@ -285,18 +382,17 @@ def parse_mlp_recompute(self, layer_idx):
pre_mlp_norm_recompute = self.mlp_rms_recompute # normalization before mlp, after attention
if self.mlp_rms_recompute:
assert self.mlp_recompute, "mlp_rms_recompute requires mlp_recompute"
- linear_recompute = self.mlp_rms_recompute or self.mlp_recompute
- router_recompute = self.mlp_rms_recompute or self.mlp_recompute
+ shared_linear_recompute = self.mlp_rms_recompute
+ linear_recompute = self.mlp_recompute
+ router_recompute = self.mlp_rms_recompute
permutation_recompute = False
else:
raise ValueError("Invalid recompute_granularity")
- enable_layer_recompute = (
- layer_idx < self.recompute_layer_num
- )
- return MLPRecomputeConfig(pre_mlp_norm_recompute = pre_mlp_norm_recompute and enable_layer_recompute,
- linear_recompute = linear_recompute and enable_layer_recompute,
- router_recompute= router_recompute and enable_layer_recompute,
- permutation_recompute = permutation_recompute and enable_layer_recompute)
+ return MLPRecomputeConfig(pre_mlp_norm_recompute = pre_mlp_norm_recompute,
+ shared_linear_recompute = shared_linear_recompute,
+ linear_recompute = linear_recompute,
+ router_recompute= router_recompute,
+ permutation_recompute = permutation_recompute)
def get_mesh_size(self, order="tp-dp-pp"):
"""According to the order to return the mesh size"""
@@ -314,29 +410,29 @@ def get_mesh_size(self, order="tp-dp-pp"):
return res
def sanity_check(self):
+ if self.cache_groupgemm_col_fp8_inputs:
+ assert self.fp8, "cache_groupgemm_col_fp8_inputs requires fp8"
+
+ if self.offload_groupgemm_col_inputs:
+ assert self.recompute_granularity != 'full_block', "offload_groupgemm_col_inputs is not allowed when recompute_granularity = 'full_block'"
+
assert (
self.world_size % self.shard_size == 0
- ), "world_size must be divisible by pp_size * tp_size"
- assert self.zero_state in [0, 1, 2, 3], "zero_state must be in [0, 1, 2, 3]"
- assert self.recompute_granularity is None or self.recompute_granularity in self.valid_recompute_granularity, f"recompute_granularity must be in [{','.join(self.valid_recompute_granularity)}]"
+ ), f"world_size must be divisible by pp_size * tp_size, but world_size = {self.world_size}, pp_size = {self.pp_size}, tp_size = {self.tp_size}"
+ assert self.zero_state in [0, 1], "zero_state must be in [0, 1]"
+ assert self.recompute_granularity is None or self.recompute_granularity in self.valid_recompute_granularity, f"recompute_granularity {self.recompute_granularity} must be in [{','.join(self.valid_recompute_granularity)}]"
assert self.recompute_layer_num >= 0
assert (
self.world_size % (self.ep_size * self.etp_size * self.pp_size) == 0
- ), "world_size must be divisible by ep_size * etp_size * pp_size"
+ ), f"world_size must be divisible by ep_size * etp_size * pp_size, but world_size = {self.world_size}, ep_size = {self.ep_size}, etp_size = {self.etp_size}, pp_size = {self.pp_size}"
assert (
self.dp_size % self.ep_size == 0
), f"dp_size {self.dp_size} is not divisible by ep_size {self.ep_size}"
- assert (
- self.ep_size == 1 or self.enable_sequence_parallel
- ), "when using ep, sp must be used"
+
assert self.moe_dispatcher_policy in [
"all2all",
"all2all-seq",
], "moe_dispatcher_policy must be in ['all2all', 'all2all-seq']"
- if self.zero_state in [2, 3]:
- assert (
- self.micro_batch_num == 1 or not self.no_sync
- ), "when using zero_state 2 and 3, no_sync must be False"
if self.interleaving_size == 1:
warnings.warn(
"interleaving_size is not supported yet, the configuration will be ignored."
@@ -351,17 +447,17 @@ def sanity_check(self):
assert (
self.etp_size == self.tp_size
), "etp_size must be equal to tp_size when using all2all-seq"
- if self.skip_ckpt_micro_batch_num > 0:
- warnings.warn(
- "skip_ckpt_micro_batch_num is not supported yet, the configuration will be ignored."
- )
- self.skip_ckpt_micro_batch_num = 0
if self.zero_state in [2, 3]:
warnings.warn(
"zero_state 2 and 3 are not supported yet, the configuration will be ignored."
)
-
+ if self.recompute_granularity == "full_block":
+ self.recompute_variance = False # megatron-LM's full recompute does not support variance
+ def reset_global_batch_size(self, global_batch_size):
+ assert global_batch_size % (self.dp_size * self.micro_batch_size)==0, f"global_batch_size {global_batch_size} must be divisible by dp_size*miro_batch_size(dp_size={self.dp_size}, micro_batch_size={self.micro_batch_size})"
+ self.micro_batch_num = global_batch_size // (self.dp_size * self.micro_batch_size)
+
@dataclass
class BandwidthConfig:
gbps: int
@@ -418,8 +514,10 @@ class SystemConfig(Config):
accelerator: AcceleratorConfig = None
networks: Dict[str, NetworkConfig] = None
real_comm_bw = {}
- FC8:bool = False
- intra_with_pcie:bool = False
+ FC8: bool = False
+ intra_with_pcie: bool = False
+ miss_efficiency: dict = field(default_factory=OrderedDict)
+ hit_efficiency: dict = field(default_factory=OrderedDict)
@classmethod
def init_from_dict(cls, config_dict: Dict[str, Any]):
@@ -454,62 +552,64 @@ def init_from_dict(cls, config_dict: Dict[str, Any]):
intra_with_pcie = intra_with_pcie,
)
- def compute_op_accuracy_time(self, op_name:str, flops:int, matmul_input_shapes:str, reture_detail=False):
+ def record_miss_efficiency(self, op_name:str, flops:int, shape_desc:str):
+ if shape_desc:
+ if op_name not in self.miss_efficiency:
+ self.miss_efficiency[op_name] = {}
+ self.miss_efficiency[op_name][f'shape={shape_desc}'] = {
+ 'flops': flops
+ }
+ def record_net_bw(self, op_name:str, net, comm_num, comm_stage:str, bw, eff_factor, latency):
+ if op_name not in self.real_comm_bw:
+ self.real_comm_bw[op_name] = {}
+ self.real_comm_bw[op_name][comm_stage.lower()] = {"net":net, "bw":f"{bw*eff_factor} GB/S", "eff_factor":eff_factor, "comm_num":comm_num, "latency": latency, "FC8":self.FC8}
+
+ def record_hit_efficiency(self, op_name:str, flops:int, shape_desc:str, eff):
+ if op_name not in self.hit_efficiency:
+ self.hit_efficiency[op_name] = {}
+ self.hit_efficiency[op_name][shape_desc] = (flops, eff)
+
+ def reset_record_info(self):
+ self.miss_efficiency.clear()
+ self.hit_efficiency.clear()
+ self.real_comm_bw.clear()
+
+ def compute_op_accuracy_time(self, op_name:str, flops:int, shape_desc:str, reture_detail=False):
"""
compute float point operation time,
return time in ms
matmul_input_shapes: list of input shapes, e.g. "[1, 16384, 4096] x [1, 4096, 128256]"
"""
+ if flops == 0:
+ if reture_detail:
+ return dict(op_name=op_name,
+ tflops=None,
+ efficient_factor=None,
+ compute_only_time = 0.0)
+ else:
+ return 0
+
op = self.accelerator.op.get(op_name, None)
if op is None:
warnings.warn(
- f"{op_name} not exist on {self.accelerator.op}, use default value"
+ f"{op_name} not exist on {self.accelerator.op.keys()}, use default value"
)
op = self.accelerator.op.get("default", None)
assert op is not None, f"default not exist on {self.accelerator.op}"
-
- if "matmul" in op_name and \
- ( op.accurate_efficient_factor is not None ) and \
- (op.accurate_efficient_factor.get(matmul_input_shapes, None) is not None):
- # matmul use accurate efficient factor to get accurate time
- efficient_factor = op.accurate_efficient_factor[matmul_input_shapes]
- print(f"matmul input shape {matmul_input_shapes} use accurate efficient factor {efficient_factor}")
- else:
- efficient_factor = op.efficient_factor
+ self.record_miss_efficiency(op_name, flops, shape_desc)
- time = flops / (op.tflops * 1e12 * efficient_factor) * 1e3
- if reture_detail:
- return dict(op_name=op_name,
- tflops=op.tflops,
- efficient_factor=efficient_factor,
- latency_us=self.accelerator.bandwidth.latency_us,
- compute_only_time = time)
- else:
- return time
-
- def compute_op_accuracy_time2(self, op_name:str, flops:int, shape_desc:str, reture_detail=False):
- """
- compute float point operation time,
- return time in ms
-
- matmul_input_shapes: list of input shapes, e.g. "[1, 16384, 4096] x [1, 4096, 128256]"
- """
- op = self.accelerator.op.get(op_name, None)
- if op is None:
- # warnings.warn(
- # f"{op_name} not exist on {self.accelerator.op.keys()}, use default value"
- # )
- op = self.accelerator.op.get("default", None)
- assert op is not None, f"default not exist on {self.accelerator.op}"
if ( op.accurate_efficient_factor is not None ) and \
(op.accurate_efficient_factor.get(shape_desc, None) is not None):
# marmul use accurate efficient factor to get accurate time
- efficient_factor = op.accurate_efficient_factor[shape_desc]
+ efficient_factor = op.accurate_efficient_factor[shape_desc]
+ self.record_hit_efficiency(op_name, flops, shape_desc, efficient_factor)
if SIMU_DEBUG:
print(f"=== \033[32m{op_name} input shape {shape_desc} use accurate compute efficient factor {efficient_factor}\033[0m, flops={flops}")
else:
efficient_factor = op.efficient_factor
+ self.record_miss_efficiency(op_name, flops, shape_desc)
+
if SIMU_DEBUG:
print(f"{op_name} input shape {shape_desc} use default compute efficient factor {efficient_factor}, flops={flops}")
@@ -522,38 +622,19 @@ def compute_op_accuracy_time2(self, op_name:str, flops:int, shape_desc:str, retu
else:
return time
- def compute_op_time(self, op_name: str, flops: int, reture_detail=False):
- """
- compute float point operation time,
- return time in ms
- """
- op = self.accelerator.op.get(op_name, None)
- if op is None:
- warnings.warn(
- f"{op_name} not exist on {self.accelerator.op}, use default value"
- )
- op = self.accelerator.op.get("default", None)
- assert op is not None, f"default not exist on {self.accelerator.op}"
- time = flops / (op.tflops * 1e12 * op.efficient_factor) * 1e3
- if reture_detail:
- return dict(op_name=op_name,
- tflops=op.tflops,
- efficient_factor=op.efficient_factor,
- compute_only_time = time)
- else:
- return time
-
def compute_mem_access_time(self, op_name, mem_bytes: int, reture_detail=False):
"""
compute memory access time,
return time in ms
"""
+
op = self.accelerator.bandwidth.get(op_name, None)
if op is None:
op = self.accelerator.bandwidth.get("default", None)
else:
if op_name != "default" and SIMU_DEBUG:
print(f'{op_name} use accurate memory bw efficiency {op.efficient_factor}')
+
time = (
mem_bytes
/ (
@@ -564,6 +645,8 @@ def compute_mem_access_time(self, op_name, mem_bytes: int, reture_detail=False):
* 1e3
)
time += op.latency_us / 1e3
+ if mem_bytes == 0:
+ time = 0
if reture_detail:
return dict(gbps=op.gbps,
efficient_factor=op.efficient_factor,
@@ -571,7 +654,7 @@ def compute_mem_access_time(self, op_name, mem_bytes: int, reture_detail=False):
io_time = time)
return time
- def compute_net_op_time(self, op_name: str, size: int, comm_num: int, net="", comm_stage="unkonw"):
+ def compute_net_op_time(self, op_name: str, size: int, comm_num: int, net="", comm_stage="unkonw", strategy:StrategyConfig=None):
"""
compute network operation time,
return time in ms
@@ -586,25 +669,37 @@ def compute_net_op_time(self, op_name: str, size: int, comm_num: int, net="", co
assert op is not None, f"{op_name} not exist on {net_data}"
scale, offset, eff_factor = op.scale, op.offset, op.efficient_factor
+ # Calculate the actual communication data based on the scale and offset of the communication operator
if eff_factor is None:
eff_factor = net_data.bandwidth.efficient_factor
actual_size = size * scale
chunk_size = actual_size / comm_num
actual_size += chunk_size * offset
+ # Specially adapted to the dp communication bandwidth of A100 pcie
if 'pcie' in net and comm_stage == 'dp' and op.dp_fixed_bw and op.dp_fixed_bw.get(str(comm_num), None):
dp_fixed_bw = op.dp_fixed_bw.get(str(comm_num))
self.real_comm_bw[op_name + "_dp"] = {"net":net, "bw":f"{dp_fixed_bw} GB/S", "comm_num":comm_num, "latency": None}
return actual_size / (dp_fixed_bw * 1024**3) * 1000
- # Bandwidth decision
+ # Intra Bandwidth decision
bw = net_data.bandwidth.gbps
- if self.FC8 and net == "high_intra_node":
+ if self.FC8 and net == "high_intra_node": # If the internal bandwidth is FC8 mode, the bandwidth changes according to the number of communications.
bw *= (comm_num-1)/7
- latency = net_data.bandwidth.latency_us
- if op.latency_us is not None:
- latency = op.latency_us
- self.real_comm_bw[op_name] = {"net":net, "bw":f"{bw*eff_factor} GB/S", "comm_num":comm_num, "latency": latency}
+
+ # Inter Bandwidth decision
+ if net == "inter_node" and strategy is not None:
+ if comm_stage == "dp":
+ # EP1, each EDP group uses 8 IBs
+ # EP2, each EDP group uses 4 IBs, ...
+ bw /= min(8, strategy.tp_size)
+ elif comm_stage == "edp":
+ # TP1, each DP group uses 8 IBs
+ # TP2, each DP group uses 4 IBs, ...
+ bw /= min(8, strategy.ep_size*strategy.etp_size)
+
+
+ latency = op.latency_us if op.latency_us is not None else net_data.bandwidth.latency_us # Default latency
if comm_num == 1:
return 0
if op in ["all_reduce", "all_gather", "reduce_scatter", "all2all"]:
@@ -613,6 +708,10 @@ def compute_net_op_time(self, op_name: str, size: int, comm_num: int, net="", co
actual_size / (bw * 1024**3 * eff_factor) * 1e3
+ (latency+fixed_latency) / 1e3
)
+ if SIMU_DEBUG:
+ if net == "high_intra_node" and op_name=="reduce_scatter":
+ print(f"op_name={op_name}, comm_num={comm_num}, net={net}, bw={bw*eff_factor} GB/S, latency={latency} us size={size}")
+ self.record_net_bw(op_name, net, comm_num, comm_stage, bw, eff_factor, latency)
return time
def compute_end2end_time(self, compute_time, mem_time):
@@ -635,40 +734,11 @@ def compute_end2end_time(self, compute_time, mem_time):
def sanity_check(self):
pass
-
- def test_gemm_time(b, m, n, k, dtype, grad_accumulation):
- """
- Test the time of gemm
- """
- try:
- import torch
- from transformer_engine.pytorch.cpp_extensions.gemm import (
- general_gemm,
- )
- except ImportError:
- raise ImportError()
- if dtype == "float16":
- dtype = torch.float16
- elif dtype == "float32":
- dtype = torch.float32
- else:
- raise NotImplementedError(f"{dtype} is not supported")
-
- if grad_accumulation:
- # grad_accumulation is not supported yet
- raise NotImplementedError("grad_accumulation is not supported yet")
-
- # test the time of gemm
- a = torch.randn(b, m, k, dtype=dtype)
- b = torch.randn(k, n, dtype=dtype)
- c = torch.randn(b, m, n, dtype=dtype)
-
@dataclass
class ModelConfig(Config):
"""Transformer model(decode-only) configuration"""
-
hidden_size: int
head_num: int
kv_head_num: int
@@ -682,7 +752,7 @@ class ModelConfig(Config):
use_swiglu: bool = None
expert_num: int = 1
topk: int = None
- attention_type: str = None
+ attention_type: str = 'mha'
moe_ffn_hidden_size: int = None
moe_shared_expert_intermediate_size: int = None
v_head_dim: int = None
@@ -691,15 +761,21 @@ class ModelConfig(Config):
q_lora_rank: int = None
kv_lora_rank: int = None
dense_layers: int = 0 # number of dense layers in moe model
- moe_pad_expert_input_to_capacity:bool = False
+ moe_pad_expert_input_to_capacity:bool = True
capacity:int = 1
group_linear_mode:str = "parallel"
make_vocab_size_divisible_by = 128 # default is 128 in megatron
padded_vocab_size = True # When tokinzer is NullTokenizer, pad vocab size to make it divisible by make_vocab_size_divisible_by * tp_size in Megatron
+
def __post_init__(self):
if self.moe_ffn_hidden_size is None:
self.moe_ffn_hidden_size = self.intermediate_size
+ if self.model_type is None:
+ if self.expert_num > 1:
+ self.model_type = 'moe'
+ else:
+ self.model_type = 'dense'
@classmethod
def init_from_config_file(cls, config_file: str):
@@ -715,8 +791,7 @@ def maybe_pad_vocab_size(self, tp_size, log=False):
Pad vocab size so it is divisible by model parallel size and
still having GPU friendly size."""
if self.padded_vocab_size:
- if self.orig_vocab_size is None:
- self.orig_vocab_size = self.vocab_size
+ self.orig_vocab_size = self.vocab_size
multiple = self.make_vocab_size_divisible_by * tp_size
after = int(math.ceil(self.orig_vocab_size / multiple) * multiple)
if log:
@@ -841,4 +916,5 @@ def layer_act_elements(self):
def sanity_check(self):
if not self.v_head_dim:
# not used for MLA
- assert self.head_num * self.head_size == self.hidden_size
+ # assert self.head_num * self.head_size == self.hidden_size
+ ...
diff --git a/simumax/core/graph.py b/simumax/core/graph.py
new file mode 100644
index 0000000..e2b6bec
--- /dev/null
+++ b/simumax/core/graph.py
@@ -0,0 +1,346 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Any, Optional, Set
+import ast
+import json
+import math
+import graphviz
+from graphviz import Digraph
+from simumax.core.tensor import FakeTensor
+from simumax.core.config import set_capture_graph_only
+
+BPE = {"fp32": 4, "fp16": 2, "bf16": 2, "fp8": 1, "int32" : 4}
+
+@dataclass
+class ValueInfo:
+ """描述张量的信息（名称、形状、数据类型）"""
+ name: str
+ shape: List[int] # 使用-1表示动态维度
+ dtype: str = "float32"
+
+@dataclass
+class Node:
+ """计算图中的节点（操作）"""
+ name: str # 节点名称
+ op: callable
+ op_type: str # 操作类型（如 'Add', 'MatMul', 'Conv'）
+ inputs: List[str] # 输入张量名称列表
+ outputs: List[str] # 输出张量名称列表
+ enable_recompute: bool = False # 是否启用重计算
+ attributes: Dict[str, Any] = None # 操作属性
+ cache_inputs: bool = False # 是否缓存输入
+ domain: str = "" # 操作域（如 '' 表示AI, 'com.microsoft' 等）
+ visited: bool = False
+
+@dataclass
+class Graph:
+ """完整的计算图"""
+ name: str = "model_graph"
+ nodes: List[Node] = None # 所有操作节点
+ inputs: List[ValueInfo] = None # 图输入
+ outputs: List[ValueInfo] = None # 图输出
+ initializers: Dict[str, Any] = None # 常量参数（权重、偏置等）
+ value_info: Dict[str, ValueInfo] = None # 中间张量信息
+ name_counts = dict() # 用于生成唯一名称的计数器
+ forward_edges: Dict[FakeTensor, List[Node]] = field(default_factory=dict) # 前向边
+ def __post_init__(self):
+ if self.nodes is None:
+ self.nodes = []
+ if self.inputs is None:
+ self.inputs = []
+ if self.outputs is None:
+ self.outputs = []
+ if self.initializers is None:
+ self.initializers = {}
+ if self.value_info is None:
+ self.value_info = {}
+
+ def reset_noede(self):
+ for node in self.nodes:
+ node.visited = False
+
+ def get_next_nodes_by_node(self, node: Node):
+ nodes = []
+ for output in node.outputs:
+ if output in self.forward_edges:
+ nodes.extend(self.forward_edges[output])
+ return nodes
+
+ def is_recompute_varaince_node(self, node: Node, next_nodes: List[Node]):
+ if not node.enable_recompute:
+ return False
+ for o in next_nodes:
+ if o.enable_recompute:
+ return False
+ return True
+
+ def traverse_forward_from_tensor(self, tensor: FakeTensor, pre_node: Node = None, set_variance_node: bool = False):
+ if isinstance(tensor, FakeTensor):
+ input_name = tensor.onnx_name
+ else:
+ input_name = tensor
+ if input_name not in self.forward_edges:
+ return
+ for node in self.forward_edges[input_name]:
+ if node.visited:
+ continue
+ node.visited = True
+ if set_variance_node and pre_node and pre_node.enable_recompute and \
+ self.is_recompute_varaince_node(node, self.get_next_nodes_by_node(node)):
+ node.op.set_variance_node(True)
+
+ for output in node.outputs:
+ self.traverse_forward_from_tensor(output, node, set_variance_node)
+
+ def to_dict(self):
+ """将图转换为字典格式，便于序列化"""
+ return {
+ "name": self.name,
+ "nodes": [{
+ "name": n.name,
+ "op_type": n.op_type,
+ "enable_recompute": n.enable_recompute,
+ "is_variance_node": n.op.is_variance_node,
+ "inputs": n.inputs,
+ "outputs": n.outputs,
+ "call_idx": n.op.call_idx,
+ "attributes": n.attributes or {},
+ "domain": n.domain
+ } for n in self.nodes],
+ "inputs": [{
+ "name": i.name,
+ "shape": i.shape,
+ "dtype": i.dtype
+ } for i in self.inputs],
+ "outputs": [{
+ "name": o.name,
+ "shape": o.shape,
+ "dtype": o.dtype
+ } for o in self.outputs],
+ "initializers": list(self.initializers.keys()),
+ "value_info": {k: {"shape": v.shape, "dtype": v.dtype}
+ for k, v in self.value_info.items()}
+ }
+
+ def export_json(self, filepath: str):
+ """导出为JSON文件"""
+ with open(filepath, 'w') as f:
+ json.dump(self.to_dict(), f, indent=2)
+
+
+class SimuONNXGraphBuilder:
+ """单例类，用于构建和存储计算图"""
+ _instance = None
+
+ def __new__(cls):
+ if cls._instance is None:
+ cls._instance = super().__new__(cls)
+ cls._instance.graph = Graph()
+ cls._instance.tensor_counter = 0
+ cls._instance.node_counter = 0
+ return cls._instance
+
+ def capture(self, perf_model):
+ print("Capture graph...")
+ set_capture_graph_only(True)
+ perf_model.run_estimate()
+ set_capture_graph_only(False)
+ print("Capture graph done.")
+ return self.graph
+
+ def reset(self):
+ """重置图构建器"""
+ self.graph = Graph()
+ self.tensor_counter = 0
+ self.node_counter = 0
+
+ def get_unique_tensor_name(self, prefix: str = "tensor") -> str:
+ """生成唯一的张量名称"""
+ name = f"{prefix}_{self.tensor_counter}"
+ self.tensor_counter += 1
+ return name
+
+ def get_unique_node_name(self, op_type: str) -> str:
+ """生成唯一的节点名称"""
+ name = f"{op_type}_{self.node_counter}"
+ self.node_counter += 1
+ # name_count = Graph.name_counts.get(op_type, 0)
+ # Graph.name_counts[op_type] = name_count + 1
+ # name = f"{op_type}_{name_count}"
+ return name
+
+ def add_node(self, op, op_type: str, inputs: List['FakeTensor'],
+ outputs: List['FakeTensor'], attributes: Dict[str, Any] = None):
+ """添加一个操作节点到图中"""
+ node_name = self.get_unique_node_name(op.name)
+ # print(f'Adding node: {op_type}, inputs={inputs}, outputs={outputs}')
+ # 确保所有输入输出张量都有名称
+ input_names = []
+ for inp in inputs:
+ if not hasattr(inp, 'onnx_name'):
+ inp.onnx_name = self.get_unique_tensor_name("input")
+ input_names.append(inp.onnx_name)
+
+ output_names = []
+ for out in outputs:
+ if not hasattr(out, 'onnx_name'):
+ out.onnx_name = self.get_unique_tensor_name("output")
+ output_names.append(out.onnx_name)
+
+ # 记录张量形状信息（简化版，实际需要更复杂的形状推断）
+ for out in outputs:
+ if hasattr(out, 'shape'):
+ self.graph.value_info[out.onnx_name] = ValueInfo(
+ name=out.onnx_name,
+ shape=out.shape if hasattr(out, 'shape') else [-1],
+ dtype=str(out.dtype) if hasattr(out, 'dtype') else "float32"
+ )
+
+ # 创建并添加节点
+ node = Node(
+ name=node_name,
+ op = op,
+ op_type=op_type,
+ inputs=input_names,
+ outputs=output_names,
+ enable_recompute=op.enable_recompute,
+ attributes=attributes or {},
+ domain=""
+ )
+ self.graph.nodes.append(node)
+ for inp in inputs:
+ name = inp.onnx_name
+ if name in self.graph.forward_edges:
+ self.graph.forward_edges[name].append(node)
+ else:
+ self.graph.forward_edges[name] = [node]
+
+ return node
+
+ @staticmethod
+ def export_json(filename):
+ SimuONNXGraphBuilder._instance.graph.export_json(filename)
+
+def export_onnx_style_graph(model, input_tensor: FakeTensor,
+ output_path: str = "model_graph.json"):
+ """
+ 导出类似ONNX的计算图
+
+ Args:
+ model: 要导出的模型
+ input_tensor: 输入张量（用于推断形状）
+ output_path: 输出JSON文件路径
+ """
+ # 重置图构建器
+ graph_builder = SimuONNXGraphBuilder()
+ graph_builder.reset()
+
+ # 设置输入张量的ONNX名称和形状信息
+ input_tensor.onnx_name = "input"
+ if hasattr(input_tensor.data, 'shape'):
+ graph_builder.graph.inputs.append(
+ ValueInfo(name=input_tensor.onnx_name,
+ shape=list(input_tensor.data.shape),
+ dtype=str(input_tensor.data.dtype))
+ )
+
+ # 执行前向传播（这会自动构建计算图）
+ output_tensor = model(input_tensor)
+
+ # 设置输出张量的信息
+ output_tensor.onnx_name = "output"
+ if hasattr(output_tensor.data, 'shape'):
+ graph_builder.graph.outputs.append(
+ ValueInfo(name=output_tensor.onnx_name,
+ shape=list(output_tensor.data.shape),
+ dtype=str(output_tensor.data.dtype))
+ )
+
+ # 添加模型参数到初始izer
+ for name, param in model._parameters.items():
+ if not hasattr(param, 'onnx_name'):
+ param.onnx_name = name
+ graph_builder.graph.initializers[param.onnx_name] = param.data
+
+ # 导出为JSON
+ graph_builder.graph.export_json(output_path)
+ print(f"计算图已导出到: {output_path}")
+
+ return graph_builder.graph
+
+def visualize_with_graphviz(json_path, output_path="computational_graph"):
+ def get_mem(size):
+ if size < 1024:
+ return f"{size} B"
+ elif size < 1024 * 1024:
+ return f"{size / 1024:.0f} KB"
+ elif size < 1024 * 1024 * 1024:
+ return f"{size / 1024 / 1024:.0f} MB"
+ else:
+ return f"{size / 1024 / 1024 / 1024:.0f} GB"
+ """使用Graphviz进行可视化"""
+ graph_data = json.load(open(json_path, 'r'))
+ # 创建有向图
+ dot = Digraph(comment=graph_data['name'])
+ dot.attr(rankdir='TB') # 从上到下布局
+ dot.attr('node', shape='rectangle', style='filled')
+
+ # 添加输入节点
+ with dot.subgraph(name='cluster_inputs') as c:
+ c.attr(color='blue', label='Inputs')
+ for inp in graph_data['inputs']:
+ c.node(inp['name'], f"Input: {inp['name']}\nShape: {inp['shape']}",
+ fillcolor='lightgreen')
+
+ # 添加输出节点
+ with dot.subgraph(name='cluster_outputs') as c:
+ c.attr(color='red', label='Outputs')
+ for out in graph_data['outputs']:
+ c.node(out['name'], f"Output: {out['name']}\nShape: {out['shape']}",
+ fillcolor='lightcoral')
+
+ # 添加参数节点
+ with dot.subgraph(name='cluster_params') as c:
+ c.attr(color='orange', label='Parameters')
+ for param in graph_data['initializers']:
+ c.node(param, f"Param: {param}", fillcolor='gold')
+
+ # 添加操作节点
+ for node in graph_data['nodes']:
+ attrs_str = '\n'.join([f"{k}: {v}" for k, v in node['attributes'].items()])
+ label = f"{node['op_type']}\n{node['name']}\ncall_idx:{node['call_idx']}"
+ if attrs_str:
+ label += f"\n---\n{attrs_str}"
+ color = 'yellow' if node['enable_recompute'] else 'lightblue'
+ color = 'green' if node['enable_recompute'] else 'lightblue'
+ # color = 'yellow' if node['is_variance_node'] else color
+ if node['is_variance_node']:
+ style = 'filled,dashed'
+ color = '#ffd700'
+ else:
+ style = 'filled,solid'
+ dot.node(node['name'], label, fillcolor=color, style=style)
+
+ # 添加中间张量节点
+ for tensor_name, info in graph_data['value_info'].items():
+ # dot.node(tensor_name, f"Tensor: {tensor_name}\nShape: {info['shape']}",
+ # fillcolor='lightgray', shape='ellipse')
+ # shape = ast.literal_eval(info['shape'])
+ shape = info['shape']
+ dtype = info['dtype']
+ mem = get_mem(math.prod(shape)*BPE[dtype])
+ dot.node(tensor_name, f"Tensor: {tensor_name}\nShape: {shape}, {dtype}\nMem:{mem}",
+ fillcolor='lightgray', shape='ellipse')
+
+ # 添加边
+ for node in graph_data['nodes']:
+ for input_tensor in node['inputs']:
+ dot.edge(input_tensor, node['name'])
+ for output_tensor in node['outputs']:
+ dot.edge(node['name'], output_tensor)
+
+ # 保存和渲染
+ dot.render(output_path, format='png', cleanup=True)
+ print(f"Graphviz图已保存到: {output_path}.png")
+
+ return dot
+
diff --git a/simumax/core/perf_llm.py b/simumax/core/perf_llm.py
index 4c135d5..9b0d898 100644
--- a/simumax/core/perf_llm.py
+++ b/simumax/core/perf_llm.py
@@ -2,6 +2,7 @@
from abc import ABC, abstractmethod
import os
+import math
import json
from copy import deepcopy
from typing import List, Union, Dict
@@ -9,9 +10,10 @@
import matplotlib.pyplot as plt
import pandas as pd
from simumax.core.base_struct import PathDebugContext
-from simumax.core.config import StrategyConfig, SystemConfig, ModelConfig, TMP_PATH, SIMU_CHECK, SIMU_DEBUG
+from simumax.core.config import StrategyConfig, SystemConfig, ModelConfig, set_capture_graph_only, TMP_PATH, SIMU_CHECK, SIMU_DEBUG, ENABLE_SIMU_GRAPH
from simumax.core.base_struct import InputOutputInfo, TensorSize, Result
from simumax.core.transformer.language_model import LLMModel, PeakPoint
+from simumax.core.graph import SimuONNXGraphBuilder, visualize_with_graphviz
from simumax.core.utils import (
HumanReadableSize,
human_readable_bytes,
@@ -20,6 +22,9 @@
rm_tmp
)
+FIRST_CHUNK = "first_stage_chunk"
+MIDDLE_CHUNK = "middle_stage_chunk"
+LAST_CHUNK = "last_stage_chunk"
class PerfBase(ABC):
"""
@@ -33,6 +38,7 @@ def __init__(self) -> None:
self.strategy = None
self.model_config = None
self.system = None
+ self.graph = None
self.debug_points = []
self.debug_points_last_stage = []
@@ -110,6 +116,7 @@ def pcie_decision_helper(size):
world_size = self.strategy.world_size
tp_size = self.strategy.tp_size
etp_size = self.strategy.etp_size
+ edp_size = self.strategy.edp_size
ep_size = self.strategy.ep_size
pp_size = self.strategy.pp_size
dp_size = self.strategy.dp_size
@@ -134,6 +141,10 @@ def pcie_decision_helper(size):
if self.strategy.dp_net == "auto" or re_analysis:
self.strategy.dp_net = pcie_decision_helper(tp_size*dp_size)
+ # 6. analysis edp_net
+ if self.strategy.edp_net == "auto" or re_analysis:
+ self.strategy.edp_net = pcie_decision_helper(etp_size * ep_size * edp_size)
+
def analysis_high_link_net(self, re_analysis):
world_size = self.strategy.world_size
tp_size = self.strategy.tp_size
@@ -153,7 +164,7 @@ def analysis_high_link_net(self, re_analysis):
# 2. analysis ep_net
if self.strategy.ep_net == "auto" or re_analysis:
- condition = (ep_size*tp_size <= num_gpu_per_nodes) # When tp *ep exceeds the number of nodes, the communication bandwidth will be reduced, and the default communication between machines will be carried out.
+ condition = (ep_size*etp_size <= num_gpu_per_nodes) # When etp *ep exceeds the number of nodes, the communication bandwidth will be reduced, and the default communication between machines will be carried out.
self.strategy.ep_net = "high_intra_node" if condition else "inter_node"
# 3. analysis tp_net
@@ -170,21 +181,39 @@ def analysis_high_link_net(self, re_analysis):
condition = (tp_size * dp_size <= num_gpu_per_nodes)
self.strategy.dp_net = "high_intra_node" if condition else "inter_node"
+ # 6. analysis edp_net
+ if self.strategy.edp_net == "auto" or re_analysis:
+ condition = etp_size * ep_size * edp_size <= num_gpu_per_nodes
+ self.strategy.edp_net = "high_intra_node" if condition else "inter_node"
+
def analysis_net(self, re_analysis = False):
if self.system.intra_with_pcie:
self.analysis_pcie_net(re_analysis)
else:
self.analysis_high_link_net(re_analysis)
-
- def run_estimate(self):
+
+ def capture(self, save_path):
+ os.makedirs(save_path, exist_ok=True)
+ print("Capture graph...")
+ builder = SimuONNXGraphBuilder()
+ builder.reset()
+ set_capture_graph_only(True)
+ self._run()
+ set_capture_graph_only(False)
+ graph = builder.graph
+ graph.export_json(os.path.join(save_path, 'model_graph.json'))
+ print("Capture graph done.")
+ return graph
+
+ def run_estimate(self, capture_graph = False, save_path='./'):
assert self.is_configured, "should call configure() first"
self.model_config.maybe_pad_vocab_size(self.strategy.tp_size)
self.analysis_net(re_analysis = True)
-
self.build()
+ if capture_graph:
+ self.graph = self.capture(save_path)
self._run()
-
class PerfLLM(PerfBase):
"""Performance model for LLM"""
@@ -324,6 +353,7 @@ def build(self):
"""
build first stage model chunk and last stage model chunk
"""
+ self.strategy.sanity_check()
self.model_chunk_dict:Dict[str, LLMModel] = {}
# Build First Stage Model Chunk
@@ -420,39 +450,6 @@ def _compute_bubble_time(self, fwd_bwd_time):
# TODO: support uneven divide && interleaving
bubble_time = fwd_bwd_time * (self.strategy.pp_size - 1)
return bubble_time
-
- def _compute_chunk_time(self, model_name='first_stage_chunk'):
- batch_stat_with_recomp = (
- self._analysis_single_batch_cost_impl(
- enable_recompute=True, model_name=model_name
- )
- )
- batch_stat_no_recomp = self._analysis_single_batch_cost_impl(
- enable_recompute=False, model_name=model_name
- )
- comm_result = self._analysis_comm_time(
- batch_stat_with_recomp, batch_stat_no_recomp, model_name=model_name
- )
- compute_result = self._analysis_compute_time(
- batch_stat_with_recomp, batch_stat_no_recomp, model_name=model_name
- )
- batch_compute_stat = compute_result[
- "batch_compute_stat"
- ]
- chunk_time = (
- batch_compute_stat["cost_info"]["fwd_compute_time"]
- + batch_compute_stat["cost_info"]["bwd_compute_time"]
- + batch_compute_stat["cost_info"]["recompute_compute_time"]
- + comm_result["intra_comm_time"][
- "intra_exposed_time_per_batch"
- ]
- + comm_result["inter_comm_time"][
- "inter_exposed_time_per_batch"
- ]
- * 2
- )
- return batch_stat_with_recomp, batch_stat_no_recomp, comm_result, compute_result, batch_compute_stat, chunk_time
-
def _compute_optim_time(self, model_name):
# we use the chunk weight accessed time as the optim time
result = {"optim_time": 0, "optim_exposed_time": 0}
@@ -500,7 +497,7 @@ def _compute_dp_time(self, model_name):
# TODO: support overlap
use_megatron = True
- def compute_dp_helper(rs_comm_size, gather_comm_size, dp_net, dp_size):
+ def compute_dp_helper(rs_comm_size, gather_comm_size, dp_net, dp_size, dp_group):
result = {"dp_comm_time": 0, "dp_comm_exposed_time": 0}
dp_comm_time = 0
bucket_size = (
@@ -518,17 +515,28 @@ def compute_dp_helper(rs_comm_size, gather_comm_size, dp_net, dp_size):
bucket_size,
comm_num=dp_size,
net=dp_net,
- comm_stage="dp"
+ comm_stage=dp_group,
+ strategy=self.strategy
)
all_gather_time = num_gather_bucket * self.system.compute_net_op_time(
- "all_gather", bucket_size, comm_num=dp_size, net=dp_net,comm_stage="dp"
+ "all_gather",
+ bucket_size,
+ comm_num=dp_size,
+ net=dp_net,
+ comm_stage=dp_group,
+ strategy=self.strategy
)
dp_comm_time += all_gather_time + reduce_scatter_time
details['reduce_scatter_time'] = reduce_scatter_time
details['all_gather_time'] = all_gather_time
else:
dp_comm_time += num_reduce_bucket * self.system.compute_net_op_time(
- "all_reduce", bucket_size, comm_num=dp_size, net=dp_net, comm_stage="dp"
+ "all_reduce",
+ bucket_size,
+ comm_num=dp_size,
+ net=dp_net,
+ comm_stage=dp_group,
+ strategy=self.strategy
)
dp_comm_exposed_time = dp_comm_time # no overlap for now
@@ -543,14 +551,16 @@ def compute_dp_helper(rs_comm_size, gather_comm_size, dp_net, dp_size):
model_info = self.model_chunk_dict[model_name].get_model_info()
- rs_comm_size = model_info.grad_bytes/2 if self.strategy.grad_reduce_in_bf16 else model_info.grad_bytes
- gather_comm_size = model_info.grad_bytes / 4 * self.dtype_to_element_size[self.strategy.dtype]
-
+ # dense
+ rs_comm_size = model_info.dense_grad_bytes/2 if self.strategy.grad_reduce_in_bf16 else model_info.dense_grad_bytes
+ gather_comm_size = model_info.dense_grad_bytes / 4 * self.dtype_to_element_size[self.strategy.dtype]
+
+ # moe
moe_rs_comm_size = model_info.moe_grad_bytes / 2 if self.strategy.grad_reduce_in_bf16 else model_info.moe_grad_bytes
moe_gather_comm_size = model_info.moe_grad_bytes / 4 * self.dtype_to_element_size[self.strategy.dtype]
- dense_dp_result = compute_dp_helper(rs_comm_size, gather_comm_size, self.strategy.dp_net, self.strategy.dp_size)
- moe_dp_result = compute_dp_helper(moe_rs_comm_size, moe_gather_comm_size, self.strategy.dp_net, self.strategy.edp_size) # FIXME: support edp_net decision
+ dense_dp_result = compute_dp_helper(rs_comm_size, gather_comm_size, self.strategy.dp_net, self.strategy.dp_size, dp_group="dp")
+ moe_dp_result = compute_dp_helper(moe_rs_comm_size, moe_gather_comm_size, self.strategy.edp_net, self.strategy.edp_size, dp_group="edp")
all_result = {
'dp_comm_exposed_time': dense_dp_result['dp_comm_exposed_time'] + moe_dp_result['dp_comm_exposed_time'],
'dense': dense_dp_result,
@@ -561,37 +571,15 @@ def compute_dp_helper(rs_comm_size, gather_comm_size, dp_net, dp_size):
def _analysis_mem_impl(
self,
micro_batch_num,
- skip_ckpt_micro_batch_num=None,
- model_name="first_stage_chunk",
+ model_name=FIRST_CHUNK,
):
- assert (
- skip_ckpt_micro_batch_num is None
- or skip_ckpt_micro_batch_num <= micro_batch_num
- )
result = {}
model_info = self.model_chunk_dict[model_name].get_model_info()
- # act_info = self.model_chunk_dict[model_name].get_act_info()
- # act_info_with_recomp = self.model_chunk_dict[
- # model_name
- # ].get_act_info_with_recomp()
- if self.strategy.enable_recompute and skip_ckpt_micro_batch_num is None:
- skip_ckpt_micro_batch_num = 0
- elif not self.strategy.enable_recompute:
- skip_ckpt_micro_batch_num = micro_batch_num
-
- cur_mb_with_recomp = micro_batch_num > skip_ckpt_micro_batch_num
- skip_ckpt_micro_batch_num_prev = (
- skip_ckpt_micro_batch_num
- if cur_mb_with_recomp
- else skip_ckpt_micro_batch_num - 1
- )
- with_recomp_micro_batch_num_prev = ( # pylint: disable=invalid-name
- (micro_batch_num - 1 - skip_ckpt_micro_batch_num)
- if cur_mb_with_recomp
- else 0
- )
+
+ #-------------------------- 0. set base info --------------------------
result["micro_batch_num"] = self.strategy.micro_batch_num
result["micro_batch_size"] = self.strategy.micro_batch_size
+ result["cached_micro_batch_num"] = micro_batch_num -1
result['parallel_config'] = {
'parallelism': self.strategy.parallelism,
'fp8': self.strategy.fp8,
@@ -600,163 +588,54 @@ def _analysis_mem_impl(
'actual_layer_num': self.model_chunk_dict['first_stage_chunk'].layer_num,
'recompute_layer': self.strategy.recompute_layer_num,
'recompute_recompute_granularity': self.strategy.recompute_granularity,
- # 'mlp_recompute_detail':self.strategy.parse_mlp_recompute(0).__dict__(),
- # 'attention_recompute_detail':self.strategy.parse_mlp_recompute(0).__dict__(),
}
}
if self.strategy.grad_reduce_in_bf16:
- model_info.grad_bytes = model_info.grad_bytes/2 # TODO(sherry): this is a hack to make it work, need to fix
+ model_info.dense_grad_bytes = model_info.dense_grad_bytes/2 # TODO(sherry): this is a hack to make it work, need to fix
+ model_info.moe_grad_bytes = model_info.moe_grad_bytes/2
+ #-------------------------- 1. compute model mem --------------------------
dense_model_mem = dict(
- model_mem = model_info.weight_bytes + model_info.grad_bytes + model_info.state_bytes,
- weight_bytes = model_info.weight_bytes,
- grad_bytes = model_info.grad_bytes,
- state_bytes = model_info.state_bytes
+ all_mem = model_info.dense_weight_bytes + model_info.dense_grad_bytes + model_info.dense_state_bytes,
+ detail = dict(
+ weight_bytes = model_info.dense_weight_bytes,
+ grad_bytes = model_info.dense_grad_bytes,
+ state_bytes = model_info.dense_state_bytes
+ )
)
moe_model_mem = dict(
- model_mem = model_info.moe_weight_bytes + model_info.moe_grad_bytes + model_info.moe_state_bytes,
- weight_bytes = model_info.moe_weight_bytes,
- grad_bytes = model_info.moe_grad_bytes,
- state_bytes = model_info.moe_state_bytes
+ all_mem = model_info.moe_weight_bytes + model_info.moe_grad_bytes + model_info.moe_state_bytes,
+ detail = dict(
+ weight_bytes = model_info.moe_weight_bytes,
+ grad_bytes = model_info.moe_grad_bytes,
+ state_bytes = model_info.moe_state_bytes
+ )
)
- result["model_mem"] = dense_model_mem["model_mem"] + moe_model_mem["model_mem"]
-
-
+ result["model_mem"] = dense_model_mem['all_mem'] + moe_model_mem['all_mem']
result["model_mem_detail"] = dict(
dense = dense_model_mem,
moe = moe_model_mem
)
- # skip gradient ckpt micro batch
- # result["act_mem_detail"] = repr(act_info)
- result["skip_ckpt_micro_batch_num_prev"] = skip_ckpt_micro_batch_num_prev
- result["with_recomp_micro_batch_num_prev"] = with_recomp_micro_batch_num_prev # micro_batch_num - 1
- result["cur_mb_with_recomp"] = cur_mb_with_recomp
-
-
- #-------------------------- 0. get with/no recompute details --------------------------
-
- cur_no_recompute_act_info:PeakPoint = self.pp_state_peak_point[model_name]["no_recompute_peak_point"]
- cur_with_recompute_act_info:PeakPoint = self.pp_state_peak_point[model_name]["with_recompute_peak_point"]
- cur_act_info:PeakPoint = cur_with_recompute_act_info if cur_mb_with_recomp else cur_no_recompute_act_info
-
- # result["cur_mbs_no_recompute_act_mem_detail"] = deepcopy(cur_no_recompute_act_info.to_dict())
- # result["cur_mbs_with_recompute_act_mem_detail"] = deepcopy(cur_with_recompute_act_info.to_dict())
- #-------------------------- get with/no recompute details --------------------------
-
-
- # result["act_cached_mem_prev_mbs"] = (
- # skip_ckpt_micro_batch_num_prev * act_info.activation_mem_cache
- # + with_recomp_micro_batch_num_prev
- # * act_info_with_recomp.activation_mem_cache
- # )
-
- #-------------------------- 1. compute act_cached_mem_prev_mbs --------------------------
- result["fwd_activation_cache_per_micro_batch"] = f"{cur_with_recompute_act_info.activation_mem_cache/1024/1024/1024:.4f} GB"
- result["peak_activation_mem_in_1F1B"] = cur_with_recompute_act_info.peak_mem
-
-
- result["act_cached_mem_prev_mbs"] = (
- skip_ckpt_micro_batch_num_prev * cur_no_recompute_act_info.activation_mem_cache
- + with_recomp_micro_batch_num_prev
- * cur_with_recompute_act_info.activation_mem_cache
- )
-
- #-------------------------- compute act_cached_mem_prev_mbs --------------------------
-
-
- # cur_act_info = act_info_with_recomp if cur_mb_with_recomp else act_info
-
- # result["cur_mbs_act_mem_detail"] = deepcopy(act_info.to_dict())
-
- peak_model_mem = result["model_mem"]
-
-
- # result["fwd_peak_allocated_mem"] = (
- # peak_model_mem
- # + result["act_cached_mem_prev_mbs"]
- # + cur_act_info.fwd_peak_mem
- # )
- # result["bwd_peak_allocated_mem"] = (
- # peak_model_mem
- # + result["act_cached_mem_prev_mbs"]
- # + cur_act_info.bwd_peak_mem
- # )
+ # result["with_recompute"] = self.strategy.enable_recompute
- #-------------------------- 2. compute fwd/bwd peak mem --------------------------
-
- result["fwd_peak_allocated_mem"] = (
- peak_model_mem
- + result["act_cached_mem_prev_mbs"]
- + cur_act_info.fwd_peak_mem
- )
-
- result["bwd_peak_allocated_mem"] = (
- peak_model_mem
- + result["act_cached_mem_prev_mbs"]
- + max(cur_act_info.bwd_peak_mem, cur_act_info.recomp_fwd_peak_mem, cur_act_info.recomp_bwd_peak_mem)
- )
- #-------------------------- compute fwd/bwd peak mem --------------------------
-
- # result["peak_cached_mem"] = (
- # max(result["bwd_peak_allocated_mem"], result["fwd_peak_allocated_mem"])
- # / self.strategy.mem_factor
- # )
+ #-------------------------- 2. compute peak activation in 1F1B--------------------------
+ cur_act_info:PeakPoint = self.pp_state_peak_point[model_name]
+ result["fwd_activation_cache_per_micro_batch"] = f"{cur_act_info.activation_mem_cache/1024/1024/1024:.4f} GB"
+ result["peak_activation_mem_in_1F1B"] = cur_act_info.peak_mem
+ model_mem = result["model_mem"]
#-------------------------- 3. compute total peak peak mem --------------------------
+ # result["fwd_peak_allocated_mem"] = cur_act_info.fwd_peak_mem
+ # result["bwd_peak_allocated_mem"] = max(cur_act_info.bwd_peak_mem, cur_act_info.recomp_fwd_peak_mem, cur_act_info.recomp_bwd_peak_mem)
result["peak_mem"] = (
- max(result["bwd_peak_allocated_mem"], result["fwd_peak_allocated_mem"])
- )
- result["peak_mem_with_reserved"] = (
- max(result["bwd_peak_allocated_mem"], result["fwd_peak_allocated_mem"])
- / self.strategy.mem_factor
+ model_mem +
+ (micro_batch_num-1) * cur_act_info.activation_mem_cache +
+ result["peak_activation_mem_in_1F1B"]
)
+ result["peak_mem_with_reserved"] = result["peak_mem"]/self.strategy.mem_factor
+
result["memory_reserved_ratio"] = str(self.strategy.mem_factor)
- result["peak_path"] = f"{cur_with_recompute_act_info.peak_path}, stage=[{cur_with_recompute_act_info.peak_stage}]"
- #-------------------------- compute total peak peak mem --------------------------
-
- is_first_stage = model_name == "first_stage_chunk"
- debug_points = (
- self.debug_points if is_first_stage else self.debug_points_last_stage
- )
- path_debug_context = (
- self.path_debug_context
- if is_first_stage
- else self.path_debug_context_last_stage
- )
-
- if debug_points:
- debug_res_info = {}
- for point in debug_points:
- data = path_debug_context.get_point_datas(
- enable_recompute=cur_mb_with_recomp
- )
- if point not in data:
- continue
- point_data = data[point]
- point_data.valid_debug_info()
- debug_res_info[point_data.point] = {
- "prev_cache_mem": human_readable_bytes(point_data.prev_cache_mem),
- "fwd_peak_no_cache_mem": human_readable_bytes(
- point_data.fwd_peak_no_cache_mem
- ),
- "bwd_peak_no_cache_mem": human_readable_bytes(
- point_data.bwd_peak_no_cache_mem
- ),
- "fwd_peak_allocated_mem": human_readable_bytes(
- peak_model_mem
- + result["act_cached_mem_prev_mbs"]
- + point_data.prev_cache_mem
- + point_data.fwd_peak_no_cache_mem
- ),
- "bwd_peak_allocated_mem": human_readable_bytes(
- peak_model_mem
- + result["act_cached_mem_prev_mbs"]
- + point_data.prev_cache_mem
- + point_data.bwd_peak_no_cache_mem
- ),
- }
- result["debug_res_info"] = debug_res_info
-
+ result["peak_path"] = f"{cur_act_info.peak_path}, stage=[{cur_act_info.peak_stage}]"
# Convert to human format
convert_final_result_to_human_format(result)
return result
@@ -765,20 +644,28 @@ def analysis_mem(self):
"""Based the simulation result, analyze the memory usage"""
if self.strategy.pp_size == 1:
result = self._analysis_mem_impl(
- micro_batch_num=1, model_name="first_stage_chunk"
+ micro_batch_num=1, model_name=FIRST_CHUNK
)
- else:
+ elif self.strategy.pp_size == 2:
+ # add more condition here to ensure the correctness the order of pp stage in result
result = {"first_stage": {}, "last_stage": {}}
result["first_stage"] = self._analysis_mem_impl(
- micro_batch_num=self.strategy.pp_size, model_name="first_stage_chunk"
- ) # 这里应该是对应的1F1B, stage1的ac需要hold pp_size份mbs（micro batch size）
+ micro_batch_num=self.strategy.pp_size, model_name=FIRST_CHUNK
+ ) # The 0th stage, here should be the corresponding 1F1B, the ac of stage1 needs to hold pp_size mbs (micro batch size)
result["last_stage"] = self._analysis_mem_impl(
- micro_batch_num=1, model_name="last_stage_chunk"
+ micro_batch_num=1, model_name=LAST_CHUNK
)
- if self.strategy.pp_size>2:
+ elif self.strategy.pp_size>2:
+ result = {"first_stage": {}, "middle_stage": {},"last_stage": {}}
+ result["first_stage"] = self._analysis_mem_impl(
+ micro_batch_num=self.strategy.pp_size, model_name=FIRST_CHUNK
+ ) # The 0th stage, here should be the corresponding 1F1B, the ac of stage1 needs to hold pp_size mbs (micro batch size)
result["middle_stage"] = self._analysis_mem_impl(
- micro_batch_num=self.strategy.pp_size-1, model_name="middle_stage_chunk"
- )# 这里应该是对应的1F1B, stage2的ac需要hold pp_size-1份mbs（micro batch size）
+ micro_batch_num=self.strategy.pp_size-1, model_name=MIDDLE_CHUNK
+ ) # The first stage, here should be the corresponding 1F1B, the ac of stage2 needs to hold pp_size-1 mbs (micro batch size)
+ result["last_stage"] = self._analysis_mem_impl(
+ micro_batch_num=1, model_name=LAST_CHUNK
+ )
return Result(result)
def _analysis_single_batch_cost_impl( # pylint: disable=invalid-name
@@ -818,7 +705,6 @@ def _analysis_single_batch_cost_impl( # pylint: disable=invalid-name
cost_batch_stat["recompute_time"] = (
cost_info.recompute_time if enable_recompute else 0
)
-
cost_batch_stat["fwd_compute_time"] = cost_info.fwd_compute_time
cost_batch_stat["bwd_compute_time"] = cost_info.bwd_compute_time
cost_batch_stat["recompute_compute_time"] = cost_info.recompute_compute_time
@@ -847,14 +733,10 @@ def _analysis_single_batch_cost_impl( # pylint: disable=invalid-name
result["compute_info"] = compute_batch_stat
return result
- def _analysis_compute_time(self, batch_stat_with_recomp, batch_stat_no_recomp, model_name):
+ def _analysis_gbs_compute_time(self, batch_stat, model_name):
result = {}
micro_batch_num = self.strategy.micro_batch_num
# skip_ckpt_micro_batch_num = self.strategy.skip_ckpt_micro_batch_num
- if self.strategy.enable_recompute:
- batch_stat = batch_stat_with_recomp
- else:
- batch_stat = batch_stat_no_recomp
result["batch_compute_stat"] = batch_stat
result["fwd_compute_time"] = (
@@ -876,17 +758,13 @@ def _analysis_compute_time(self, batch_stat_with_recomp, batch_stat_no_recomp, m
result["model_flops"] = result["fwd_flops"] + result["bwd_flops"]
return result
- def _analysis_comm_time(self, batch_stat_with_recomp, batch_stat_no_recomp, model_name):
+ def _analysis_gbs_comm_time(self, batch_stat, model_name):
result = {}
micro_batch_num = self.strategy.micro_batch_num
dp_comm_result = self._compute_dp_time(model_name)
# TODO: add ckpt bubble and add strategy extra comm time, # e.g sp grad reduce
- intra_exposed_time_with_recomp_per_batch = sum( # pylint: disable=invalid-name
- batch_stat_with_recomp["cost_info"][k]
- for k in ["fwd_net_time", "bwd_net_time", "recompute_net_time"]
- )
- intra_exposed_time_no_recomp_per_batch = sum( # pylint: disable=invalid-name
- batch_stat_no_recomp["cost_info"][k]
+ intra_exposed_time = sum( # pylint: disable=invalid-name
+ batch_stat["cost_info"][k]
for k in ["fwd_net_time", "bwd_net_time", "recompute_net_time"]
)
if self.strategy.pp_size > 1:
@@ -899,29 +777,18 @@ def _analysis_comm_time(self, batch_stat_with_recomp, batch_stat_no_recomp, mode
if self.strategy.enable_sequence_parallel
else pp_comm_size
)
- inter_exposed_time_per_batch = 2 * self.system.compute_net_op_time(
+ inter_exposed_time_per_batch = 2 * 2 * self.system.compute_net_op_time(
"p2p", pp_comm_size, 2, net=self.strategy.pp_net, comm_stage="pp"
- ) # 2 p2p
-
+ ) # 2 p2p, 2 to fwd and bwd
else:
inter_exposed_time_per_batch = 0
- # intra_exposed_time_with_recomp = intra_exposed_time_with_recomp_per_batch
- # intra_exposed_time_with_recomp = intra_exposed_time_with_recomp_per_batch * \
- # micro_batch_num
- # intra_exposed_time_no_recomp = intra_exposed_time_no_recomp_per_batch * micro_batch_num
+
inter_exposed_time = inter_exposed_time_per_batch * micro_batch_num
result["dp_comm_time"] = dp_comm_result
# Now we don't consider the mix of recompute and non-recompute
- if self.strategy.enable_recompute:
- intra_exposed_time_per_batch = intra_exposed_time_with_recomp_per_batch
- intra_exposed_time = (
- intra_exposed_time_with_recomp_per_batch * micro_batch_num
- )
- else:
- intra_exposed_time_per_batch = intra_exposed_time_no_recomp_per_batch
- intra_exposed_time = (
- intra_exposed_time_no_recomp_per_batch * micro_batch_num
- )
+ intra_exposed_time_per_batch = intra_exposed_time
+ intra_exposed_time = intra_exposed_time_per_batch * micro_batch_num
+
result["intra_comm_time"] = {
"intra_exposed_time_per_batch": intra_exposed_time_per_batch,
"intra_exposed_time": intra_exposed_time,
@@ -982,7 +849,6 @@ def calculate_1f1b_bubble(self, pp, mbc, forward_times, backward_times, draw=Fal
# idx = np.argmax(f_b_time)
# bubble = sum(f_b_time)+sum(forward_times[idx+1:])+sum(backward_times[idx+1:]) - (pp-idx)*(f_b_time[idx])
# all_time = bubble + mbc*f_b_time[idx]
-
if draw:
# 可视化调度图
fig, ax = plt.subplots(figsize=(12, 5))
@@ -1005,224 +871,53 @@ def calculate_1f1b_bubble(self, pp, mbc, forward_times, backward_times, draw=Fal
plt.savefig("corrected_1F1B_pipeline.png")
return max_time
-
- def _analysis_single_iter_cost_impl2(self):
- FIRST_CHUNK = "first_stage_chunk"
- MIDDLE_CHUNK = "middle_stage_chunk"
- LAST_CHUNK = "last_stage_chunk"
-
- def compute_single_iter_cost_helper(model_name = "first_stage_chunk"):
- assert model_name in ["first_stage_chunk", "middle_stage_chunk", 'last_stage_chunk'], f"model_name {model_name} not supported"
- batch_stat_with_recomp = self._analysis_single_batch_cost_impl(
- enable_recompute=True, model_name=model_name
- )
- batch_stat_no_recomp = self._analysis_single_batch_cost_impl(
- enable_recompute=False, model_name=model_name
- )
- # 1.comm_result： dp_time + fwd/bwd/recompute net time + pp_time
- comm_cost_result = self._analysis_comm_time(
- batch_stat_with_recomp, batch_stat_no_recomp, model_name
- )
- # 2.compute result:
- # - fwd/bwd/recompute compute time, stored in compute_result['cost_info'], summarized in compute_result['fwd/bwd/recompute_compute_time']. compute_result['cost_info']['dp_comm_time] is computed by model.grad_bytes.
- # - optim update time, stored in compute_result['optim_time'], computed by model.state_bytes.
- # - fwd/bwd/recompute fwd/bwd/recompute net time is ignored in compute time, but included in comm_result['intra_comm_time']
- compute_cost_result = self._analysis_compute_time(
- batch_stat_with_recomp, batch_stat_no_recomp
- )
-
- breakdown_result = {}
- breakdown_result["fwd_compute_time"] = compute_cost_result["fwd_compute_time"]
- breakdown_result["recompute_time"] = compute_cost_result["recompute_time"]
- breakdown_result["bwd_compute_time"] = compute_cost_result["bwd_compute_time"]
- breakdown_result["optim_time"] = compute_cost_result["optim_time"][
- "optim_exposed_time"
- ]
- # breakdown_result['grad_accumulation'] = self.strategy.micro_batch_num * self.system.compute_mem_access_time(self.model_chunk_dict[model_name]._model_info.grad_bytes)
- breakdown_result["intra_exposed_time"] = comm_cost_result["intra_comm_time"][
- "intra_exposed_time"
- ]
- breakdown_result["inter_exposed_time"] = comm_cost_result["inter_comm_time"][
- "inter_exposed_time"
- ]
- breakdown_result["dp_exposed_time"] = comm_cost_result["dp_comm_time"][
- "dp_comm_exposed_time"
- ]
- breakdown_result['flexible_time'] = 20 # traning log/h2d
-
- cost_info = self.model_chunk_dict[model_name]._cost_info
- chunk_time = (
- cost_info.fwd_compute_time
- + cost_info.bwd_compute_time
- + cost_info.recompute_compute_time
- + comm_cost_result["intra_comm_time"]["intra_exposed_time_per_batch"]
- + comm_cost_result["inter_comm_time"]["inter_exposed_time_per_batch"] * 2
- )
- breakdown_result['bubble_time'] = self._compute_bubble_time(chunk_time)
-
- all_tokens_per_iter = self.strategy.seq_len * self.strategy.global_batch_size
-
- all_result = {}
- all_result["comm_cost_details"] = comm_cost_result
- all_result["compute_cost_details"] = compute_cost_result
- all_result["breakdown_result"] = breakdown_result
- all_result['micro_batch_num'] = self.strategy.micro_batch_num
- # all_result["1F1B_time(=breakdown_result_time remove dp/optim/bubble time)"] = chunk_time
- all_result["all_tokens_per_iter"] = all_tokens_per_iter
-
- return all_result
-
- # def compute_pp_total_time_helper():
- # def compute_fwd_bwd_time(model_name):
- # # if len(compute_result) == 0:
- # # return -1, -1
- # # chunk_time = (
- # # batch_compute_stat["cost_info"]["fwd_compute_time"]
- # # + batch_compute_stat["cost_info"]["bwd_compute_time"]
- # # + batch_compute_stat["cost_info"]["recompute_compute_time"]
- # # + comm_result["intra_comm_time"]["intra_exposed_time_per_batch"]
- # # + comm_result["inter_comm_time"]["inter_exposed_time_per_batch"] * 2
- # # )
- # if self.strategy.pp_size > 1:
- # pp_comm_size = (
- # self.micro_hidden_states_size
- # * self.dtype_to_element_size[self.strategy.dtype]
- # )
- # pp_comm_size = (
- # pp_comm_size / self.strategy.tp_size
- # if self.strategy.enable_sequence_parallel
- # else pp_comm_size
- # )
- # pp_time = 2 * self.system.compute_net_op_time(
- # "p2p", pp_comm_size, 2, net=self.strategy.pp_net
- # ) # 2 p2p
-
- # else:
- # pp_time = 0
-
- # cost_info = self.model_chunk_dict[model_name].get_cost_info()
-
- # fwd_chunk_time = (cost_info.fwd_compute_time +
- # cost_info.fwd_net_time +
- # pp_time)
-
- # bwd_chunk_time = (cost_info.bwd_compute_time +
- # cost_info.bwd_net_time +
- # cost_info.recompute_compute_time +
- # cost_info.recompute_net_time +
- # pp_time)
- # # fwd_time = (cost_info['fwd_compute_time'] +
- # # cost_info['fwd_net_exposed_time']
- # # )
- # # bwd_time = (cost_info['bwd_compute_time'] +
- # # cost_info['bwd_net_exposed_time'] +
- # # cost_info['recompute_compute_time'] +
- # # cost_info['recompute_net_exposed_time']
- # # )
- # return fwd_chunk_time, bwd_chunk_time
- # fwd_chunk_time, bwd_chunk_time = compute_fwd_bwd_time(FIRST_CHUNK)
- # forward_times = [fwd_chunk_time]
- # backward_times = [bwd_chunk_time]
- # has_middle_chunks = self.strategy.pp_size > 2
- # has_last_chunk = self.strategy.pp_size > 1
- # if has_middle_chunks:
- # fwd_chunk_time, bwd_chunk_time = compute_fwd_bwd_time(MIDDLE_CHUNK)
- # forward_times.extend([fwd_chunk_time]*(self.strategy.pp_size - 2))
- # backward_times.extend([bwd_chunk_time]*(self.strategy.pp_size - 2))
- # if has_last_chunk:
- # fwd_chunk_time, bwd_chunk_time = compute_fwd_bwd_time(LAST_CHUNK)
- # forward_times.append(fwd_chunk_time)
- # backward_times.append(bwd_chunk_time)
-
- # single_iter_time = self.calculate_1f1b_bubble(self.strategy.pp_size, self.strategy.micro_batch_num, forward_times, backward_times)
- # single_iter_time += self._compute_dp_time()['dp_comm_exposed_time']
- # single_iter_time += self._compute_optim_time()['optim_exposed_time']
- # return single_iter_time
-
- def compute_mfu_helper(all_result, duration_time_per_iter):
- # breakdown_result = all_result['breakdown_result']
- # breakdown_result['bubble_time'] = bubble_time
- # breakdown_result['uneven_bubble_time'] = uneven_bubble_time
- duration_time_per_iter = duration_time_per_iter
- # all_result['bubble_time'] = bubble_time
- # all_result['uneven_bubble_time'] = uneven_bubble_time
- all_result['duration_time_per_iter'] = duration_time_per_iter
- # TODO(sherry): add bubble time
-
- chunk_time = sum(all_result['breakdown_result'].values())
- all_result['breakdown_result']['bubble_time'] = duration_time_per_iter - chunk_time
-
- throughput_per_accelerator = (
- all_result["all_tokens_per_iter"]
- / (duration_time_per_iter / 1000)
- / self.strategy.world_size
- )
- all_result["throughput_per_accelerator"] = throughput_per_accelerator
- all_result["mfu_6nd_with_attn"] = (
- self.model_config.flops_per_token(
- context_seq_len=self.strategy.seq_len, with_attn=True
+ def _compute_single_batch_fwd_bwd_time(self, model_name, chunk = False):
+ if self.strategy.pp_size > 1:
+ pp_comm_size = (
+ self.micro_hidden_states_size
+ * self.dtype_to_element_size[self.strategy.dtype]
)
- * throughput_per_accelerator
- / self.system.accelerator.op["default"].tflops
- / 1e12 # convert byte-flops to tflops
- )
-
- compute_cost_result = all_result['compute_cost_details']
- mfu = (
- compute_cost_result["model_flops"]
- / (duration_time_per_iter / 1000) # convert ms to s
- / self.system.accelerator.op["default"].tflops
- / 1e12 # convert byte-flops to tflops
- )
- TFLOPs = (compute_cost_result["model_flops"]
- / (duration_time_per_iter / 1000) # convert ms to s
- / 1e12
+ pp_comm_size = (
+ pp_comm_size / self.strategy.tp_size
+ if self.strategy.enable_sequence_parallel
+ else pp_comm_size
)
-
- all_result["mfu"] = mfu
- all_result['model_flops'] = compute_cost_result["model_flops"]
- all_result["TFLOPs"] = TFLOPs
- return mfu
-
- res = {
- FIRST_CHUNK: {},
- MIDDLE_CHUNK: {},
- LAST_CHUNK: {},
- }
-
-
- res[FIRST_CHUNK] = compute_single_iter_cost_helper(model_name=FIRST_CHUNK)
- if self.strategy.pp_size > 2:
- res[MIDDLE_CHUNK] = compute_single_iter_cost_helper(model_name=MIDDLE_CHUNK)
-
- if self.strategy.pp_size > 1:
- res[LAST_CHUNK] = compute_single_iter_cost_helper(model_name=LAST_CHUNK)
-
- single_iter_time = sum(v for _, v in res[FIRST_CHUNK]['breakdown_result'].items())
-
- mfu = 0
- TFLOPs = 0
- mfu += compute_mfu_helper(res[FIRST_CHUNK], single_iter_time)
- TFLOPs += res[FIRST_CHUNK]["TFLOPs"]
- if self.strategy.pp_size > 1:
- mfu += compute_mfu_helper(res[LAST_CHUNK], single_iter_time)
- TFLOPs += res[LAST_CHUNK]["TFLOPs"]
- if self.strategy.pp_size > 2:
- mfu += compute_mfu_helper(res[MIDDLE_CHUNK], single_iter_time) * (self.strategy.pp_size -2)
- TFLOPs += res[MIDDLE_CHUNK]["TFLOPs"] * (self.strategy.pp_size -2)
-
-
- mfu /= self.strategy.pp_size
- TFLOPs /= self.strategy.pp_size
-
- res['mfu'] = mfu
- res['mfu_6nd_with_attn'] = res[FIRST_CHUNK]['mfu_6nd_with_attn']
- res['TFLOPs'] = TFLOPs
- res['duration_time_per_iter'] = single_iter_time
-
- convert_final_result_to_human_format(res)
- return res
+ pp_time = 2 * self.system.compute_net_op_time(
+ "p2p", pp_comm_size, 2, net=self.strategy.pp_net
+ ) # 2 p2p, fwd/bwd each
+ else:
+ pp_time = 0
+
+ cost_info = self.model_chunk_dict[model_name].get_cost_info()
+
+ fwd_chunk_time = (cost_info.fwd_compute_time +
+ cost_info.fwd_net_time +
+ pp_time)
+ bwd_chunk_time = (cost_info.bwd_compute_time +
+ cost_info.bwd_net_time +
+ cost_info.recompute_compute_time +
+ cost_info.recompute_net_time +
+ pp_time)
+ return (fwd_chunk_time, bwd_chunk_time) if not chunk else fwd_chunk_time + bwd_chunk_time
+
+ def _compute_pp_total_time(self):
+ fwd_chunk_time, bwd_chunk_time = self._compute_single_batch_fwd_bwd_time(FIRST_CHUNK)
+ forward_times = [fwd_chunk_time]
+ backward_times = [bwd_chunk_time]
+ has_middle_chunks = self.strategy.pp_size > 2
+ has_last_chunk = self.strategy.pp_size > 1
+ if has_middle_chunks:
+ fwd_chunk_time, bwd_chunk_time = self._compute_single_batch_fwd_bwd_time(MIDDLE_CHUNK)
+ forward_times.extend([fwd_chunk_time]*(self.strategy.pp_size - 2))
+ backward_times.extend([bwd_chunk_time]*(self.strategy.pp_size - 2))
+ if has_last_chunk:
+ fwd_chunk_time, bwd_chunk_time = self._compute_single_batch_fwd_bwd_time(LAST_CHUNK)
+ forward_times.append(fwd_chunk_time)
+ backward_times.append(bwd_chunk_time)
+
+ single_iter_time = self.calculate_1f1b_bubble(self.strategy.pp_size, self.strategy.micro_batch_num, forward_times, backward_times, draw = False)
+ return single_iter_time
def _analysis_single_iter_cost_impl(self):
# we construct the result in the following hierarchy:
@@ -1231,213 +926,156 @@ def _analysis_single_iter_cost_impl(self):
# third level-0: compute time = fwd time + recom_time + bwd_time + optim update time
# third level-1: comm_time_: tp_time(tp_time、tp_time_can_overlap) + pp_time
all_result = {}
- batch_stat_with_recomp = self._analysis_single_batch_cost_impl(
- enable_recompute=True, model_name = "first_stage_chunk"
- )
- batch_stat_no_recomp = self._analysis_single_batch_cost_impl(
- enable_recompute=False, model_name = "first_stage_chunk"
+ single_batch_cost = self._analysis_single_batch_cost_impl(
+ enable_recompute=self.strategy.enable_recompute, model_name = FIRST_CHUNK
)
# 1.comm_result： dp_time + fwd/bwd/recompute net time + pp_time
- comm_result = self._analysis_comm_time(
- batch_stat_with_recomp, batch_stat_no_recomp, model_name = "first_stage_chunk"
- )
+ gbs_comm_in_first_stage = self._analysis_gbs_comm_time(single_batch_cost, model_name = FIRST_CHUNK)
# 2.compute result:
- # - fwd/bwd/recompute compute time, stored in compute_result['cost_info'], summarized in compute_result['fwd/bwd/recompute_compute_time']. compute_result['cost_info']['dp_comm_time] is computed by model.grad_bytes.
- # - optim update time, stored in compute_result['optim_time'], computed by model.state_bytes.
- # - fwd/bwd/recompute fwd/bwd/recompute net time is ignored in compute time, but included in comm_result['intra_comm_time']
- compute_result_first_stage = self._analysis_compute_time(
- batch_stat_with_recomp, batch_stat_no_recomp, model_name = "first_stage_chunk"
+ gbs_compute_cost_in_first_stage = self._analysis_gbs_compute_time(
+ single_batch_cost, model_name = FIRST_CHUNK
)
# 3. all time
- batch_compute_stat = compute_result_first_stage["batch_compute_stat"]
- bubble_uneven_time = 0
-
# can't be overlap for now
- # first_stage chunk time
- chunk_time = (
- batch_compute_stat["cost_info"]["fwd_compute_time"]
- + batch_compute_stat["cost_info"]["bwd_compute_time"]
- + batch_compute_stat["cost_info"]["recompute_compute_time"]
- + comm_result["intra_comm_time"]["intra_exposed_time_per_batch"]
- + comm_result["inter_comm_time"]["inter_exposed_time_per_batch"] * 2
- )
+ chunk_time = self._compute_single_batch_fwd_bwd_time(FIRST_CHUNK, chunk=True)
if self.strategy.pp_size > 1:
- batch_stat_with_recomp_last_stage = ( # pylint: disable=invalid-name
- self._analysis_single_batch_cost_impl(
- enable_recompute=True, model_name="last_stage_chunk"
+ single_batch_cost = self._analysis_single_batch_cost_impl(
+ enable_recompute=self.strategy.enable_recompute, model_name=LAST_CHUNK
)
+ gbs_comm_result_in_last_stage = self._analysis_gbs_comm_time(
+ single_batch_cost, model_name=LAST_CHUNK
)
- batch_stat_no_recomp_last_stage = self._analysis_single_batch_cost_impl(
- enable_recompute=False, model_name="last_stage_chunk"
+ gbs_compute_result_in_last_stage = self._analysis_gbs_compute_time(
+ single_batch_cost, model_name=LAST_CHUNK
)
- comm_result_last_stage = self._analysis_comm_time(
- batch_stat_with_recomp_last_stage, batch_stat_no_recomp_last_stage, model_name="last_stage_chunk"
- )
- compute_result_last_stage = self._analysis_compute_time(
- batch_stat_with_recomp_last_stage, batch_stat_no_recomp_last_stage, model_name="last_stage_chunk"
- )
- batch_compute_stat_last_stage = compute_result_last_stage[
- "batch_compute_stat"
- ]
- chunk_time_last_stage = (
- batch_compute_stat_last_stage["cost_info"]["fwd_compute_time"]
- + batch_compute_stat_last_stage["cost_info"]["bwd_compute_time"]
- + batch_compute_stat_last_stage["cost_info"]["recompute_compute_time"]
- + comm_result_last_stage["intra_comm_time"][
- "intra_exposed_time_per_batch"
- ]
- + comm_result_last_stage["inter_comm_time"][
- "inter_exposed_time_per_batch"
- ]
- * 2
- )
- bubble_uneven_time = (
- abs((chunk_time_last_stage - chunk_time))
- * self.strategy.micro_batch_num
- )
-
- bubble_time = self._compute_bubble_time(chunk_time) # bubble_time = chunk_time * (self.strategy.pp_size - 1)
+ chunk_time_lstage = self._compute_single_batch_fwd_bwd_time(LAST_CHUNK, chunk=True)
breakdown_result = {}
- breakdown_result["fwd_compute_time"] = compute_result_first_stage["fwd_compute_time"]
- breakdown_result["recompute_time"] = compute_result_first_stage["recompute_time"]
- breakdown_result["bwd_compute_time"] = compute_result_first_stage["bwd_compute_time"]
- breakdown_result["optim_time"] = compute_result_first_stage["optim_time"][
+ breakdown_result["fwd_compute_time"] = gbs_compute_cost_in_first_stage["fwd_compute_time"]
+ breakdown_result["recompute_time"] = gbs_compute_cost_in_first_stage["recompute_time"]
+ breakdown_result["bwd_compute_time"] = gbs_compute_cost_in_first_stage["bwd_compute_time"]
+ breakdown_result["optim_time"] = gbs_compute_cost_in_first_stage["optim_time"][
"optim_exposed_time"
]
- breakdown_result["intra_exposed_time"] = comm_result["intra_comm_time"][
+ breakdown_result["intra_exposed_time"] = gbs_comm_in_first_stage["intra_comm_time"][
"intra_exposed_time"
]
- breakdown_result["inter_exposed_time"] = comm_result["inter_comm_time"][
+ breakdown_result["inter_exposed_time"] = gbs_comm_in_first_stage["inter_comm_time"][
"inter_exposed_time"
]
- breakdown_result["dp_exposed_time"] = comm_result["dp_comm_time"][
+ breakdown_result["dp_exposed_time"] = gbs_comm_in_first_stage["dp_comm_time"][
"dp_comm_exposed_time"
]
- breakdown_result["bubble_time"] = bubble_time
- breakdown_result["bubble_uneven_time"] = bubble_uneven_time
- all_time = sum(v for _, v in breakdown_result.items())
+
if self.strategy.pp_size > 1:
breakdown_result_last_stage = {}
- breakdown_result_last_stage["fwd_compute_time"] = compute_result_last_stage["fwd_compute_time"]
- breakdown_result_last_stage["recompute_time"] = compute_result_last_stage["recompute_time"]
- breakdown_result_last_stage["bwd_compute_time"] = compute_result_last_stage["bwd_compute_time"]
- breakdown_result_last_stage["optim_time"] = compute_result_last_stage["optim_time"][
+ breakdown_result_last_stage["fwd_compute_time"] = gbs_compute_result_in_last_stage["fwd_compute_time"]
+ breakdown_result_last_stage["recompute_time"] = gbs_compute_result_in_last_stage["recompute_time"]
+ breakdown_result_last_stage["bwd_compute_time"] = gbs_compute_result_in_last_stage["bwd_compute_time"]
+ breakdown_result_last_stage["optim_time"] = gbs_compute_result_in_last_stage["optim_time"][
"optim_exposed_time"
]
- breakdown_result_last_stage["intra_exposed_time"] = comm_result_last_stage["intra_comm_time"][
+ breakdown_result_last_stage["intra_exposed_time"] = gbs_comm_result_in_last_stage["intra_comm_time"][
"intra_exposed_time"
]
- breakdown_result_last_stage["inter_exposed_time"] = comm_result_last_stage["inter_comm_time"][
+ breakdown_result_last_stage["inter_exposed_time"] = gbs_comm_result_in_last_stage["inter_comm_time"][
"inter_exposed_time"
]
- breakdown_result_last_stage["dp_exposed_time"] = comm_result_last_stage["dp_comm_time"][
+ breakdown_result_last_stage["dp_exposed_time"] = gbs_comm_result_in_last_stage["dp_comm_time"][
"dp_comm_exposed_time"
]
if self.strategy.pp_size > 2:
- chunk_time_middle_stage = self._compute_chunk_time(model_name='middle_stage_chunk')[-1]
+ chunk_time_middle_stage = self._compute_single_batch_fwd_bwd_time(MIDDLE_CHUNK, chunk=True)
else:
chunk_time_middle_stage = 0
-
- bubble_time_last_stage = chunk_time + chunk_time_middle_stage*(self.strategy.pp_size-2)
- breakdown_result_last_stage["bubble_time"] = bubble_time_last_stage
- all_time_last_stage = sum(v for _, v in breakdown_result_last_stage.items())
- breakdown_result_last_stage['all_time'] = all_time_last_stage
all_result["breakdown_result_last_stage"] = breakdown_result_last_stage
- # all_time = all_time_last_stage # FIXME(sherry)：使用last stage的all_time
+
# 4.compute first level
- all_tokens_per_iter = self.strategy.seq_len * self.strategy.global_batch_size
- duration_time_per_iter = all_time # TODO(sherry)：使用all time进行后续所有stage的mfu计算
- model_flops = compute_result_first_stage["model_flops"]
- throughput_per_accelerator = (
- all_tokens_per_iter
- / (duration_time_per_iter / 1000)
- / self.strategy.world_size
- )
- mfu = (
- model_flops
- / (duration_time_per_iter / 1000) # convert ms to s
- / self.system.accelerator.op["default"].tflops
- / 1e12 # convert byte-flops to tflops
- )
- TFLOPS = (model_flops
- / (duration_time_per_iter / 1000) # convert ms to s
- / 1e12
- )
+ model_flops = gbs_compute_cost_in_first_stage["model_flops"]
+ # ------------------------- SUMMRY -------------------------
+ pp_size = self.strategy.pp_size
+ dense_param_numel = self.model_chunk_dict[FIRST_CHUNK]._model_info.weight_numel + (
+ self.model_chunk_dict[MIDDLE_CHUNK]._model_info.weight_numel if pp_size > 2 else 0
+ ) * (pp_size - 2) + (
+ self.model_chunk_dict[LAST_CHUNK]._model_info.weight_numel if pp_size > 1 else 0
+ )
+ moe_param_numel = self.model_chunk_dict[FIRST_CHUNK]._model_info.moe_weight_numel + (
+ self.model_chunk_dict[MIDDLE_CHUNK]._model_info.moe_weight_numel if pp_size > 2 else 0
+ ) * (pp_size - 2) + (
+ self.model_chunk_dict[LAST_CHUNK]._model_info.moe_weight_numel if pp_size > 1 else 0
+ )
+
+ def get_dp_and_optim(model_chunk):
+ t = self._compute_dp_time(model_chunk)['dp_comm_exposed_time']
+ t += self._compute_optim_time(model_chunk)['optim_exposed_time']
+ return t
+ single_iter_time_no_dp_opim = self._compute_pp_total_time()
+ duration_times = [single_iter_time_no_dp_opim + get_dp_and_optim(FIRST_CHUNK)]
+ duration_times.append(single_iter_time_no_dp_opim + get_dp_and_optim(MIDDLE_CHUNK)) if self.strategy.pp_size > 2 else 0
+ duration_times.append(single_iter_time_no_dp_opim + get_dp_and_optim(LAST_CHUNK)) if self.strategy.pp_size > 1 else 0
+
+ final_duration_time_per_iter = max(duration_times)
+ all_tokens_per_iter = self.strategy.seq_len * self.strategy.global_batch_size
+
+ theory_flops_per_token = self.model_config.flops_per_token(context_seq_len=self.strategy.seq_len, with_attn=True)
theory_flops = self.model_config.flops_per_token(context_seq_len=self.strategy.seq_len, with_attn=True) * all_tokens_per_iter // self.strategy.world_size
- all_result["comm_details"] = comm_result
- all_result["compute_details"] = compute_result_first_stage
+ TGS = all_tokens_per_iter/(final_duration_time_per_iter/1000)/self.strategy.world_size
+ TFLOPS = theory_flops / (final_duration_time_per_iter/1000)/1e12
+ TFLOPS_PER_TOKEN = theory_flops_per_token / (final_duration_time_per_iter/1000)/1e12
+ new_mfu_6nd_with_attn = TFLOPS / self.system.accelerator.op["default"].tflops
+
+ mbc = self.strategy.micro_batch_num
+ all_result["comm_details"] = gbs_comm_in_first_stage
+ all_result["compute_details"] = gbs_compute_cost_in_first_stage
all_result["breakdown_result"] = breakdown_result
- all_result["chunk_time"] = chunk_time
all_result["all_tokens_per_iter"] = all_tokens_per_iter
- all_result["duration_time_per_iter"] = duration_time_per_iter
- all_result["mfu_6nd_with_attn"] = (
- self.model_config.flops_per_token(
- context_seq_len=self.strategy.seq_len, with_attn=True
- )
- * throughput_per_accelerator
- / self.system.accelerator.op["default"].tflops
- / 1e12 # convert byte-flops to tflops
- )
- if self.strategy.pp_size > 1:
- mfu_last_stage = (
- compute_result_last_stage["model_flops"]
- / (duration_time_per_iter / 1000)
- / self.system.accelerator.op["default"].tflops
- / 1e12
- )
- mfu += mfu_last_stage
- tflops_last_stage =(
- compute_result_last_stage["model_flops"]
- / (duration_time_per_iter / 1000)
- / 1e12
- )
- TFLOPS += tflops_last_stage
- if self.strategy.pp_size > 2:
- batch_stat_with_recomp_middle_stage = (
- self._analysis_single_batch_cost_impl(
- enable_recompute=True, model_name="middle_stage_chunk"
- )
- )
- batch_stat_no_recomp_middle_stage = self._analysis_single_batch_cost_impl(
- enable_recompute=False, model_name="middle_stage_chunk"
- )
- compute_result_middle_stage = self._analysis_compute_time(
- batch_stat_with_recomp_middle_stage, batch_stat_no_recomp_middle_stage, model_name="middle_stage_chunk"
- )
- mfu_middle_stage = (
- compute_result_middle_stage["model_flops"]
- / (duration_time_per_iter / 1000)
- / self.system.accelerator.op["default"].tflops
- / 1e12
- )
- mfu += mfu_middle_stage*(self.strategy.pp_size-2)
- tflops_middle_stage = (
- compute_result_middle_stage["model_flops"]
- / (duration_time_per_iter / 1000)
- / 1e12
- )
- TFLOPS += tflops_middle_stage * (self.strategy.pp_size - 2)
- mfu /= self.strategy.pp_size
- TFLOPS /= self.strategy.pp_size
- # mfu_6nd_with_attn: The MFU is calculated based on the FLOPS of standard attention
- # mfu: The MFU is calculated based on the actual FLOPS of attention. If it is flashattention, it will have one more base flops than standard attention
- all_result["mfu"] = mfu
-
- all_result["throughput_per_accelerator"] = throughput_per_accelerator
- all_result["throughput per GPU (TFLOP/s/GPU)"] = theory_flops / (duration_time_per_iter/1000)/1e12
+
+ def format_chunk_time(model_chunk, chunk_time, duration_time):
+ return {
+ model_chunk:{
+ 'duration_time(chunk_timexmbc+dp_optim+bubble)': duration_time,
+ 'chunk_time(fwd+bwd)':chunk_time,
+ 'dp_and_optim_time': get_dp_and_optim(model_chunk),
+ 'bubble_time': single_iter_time_no_dp_opim - mbc*(chunk_time)
+ }
+ }
+ all_result['all_chunk_times'] = format_chunk_time(FIRST_CHUNK, chunk_time, duration_times[0])
+ all_result['all_chunk_times'].update(format_chunk_time(MIDDLE_CHUNK, chunk_time_middle_stage, duration_times[1]) if pp_size > 2 else {})
+ all_result['all_chunk_times'].update(format_chunk_time(LAST_CHUNK, chunk_time_lstage, duration_times[-1]) if pp_size > 1 else {})
+
+ all_result.update({
+ 'duration_time_per_iter': final_duration_time_per_iter,
+ 'throughput_per_accelerator': TGS,
+ 'throughput per GPU (TFLOP/s/GPU)': TFLOPS,
+ 'throughput per GPU per token (TFLOP/s/GPU/token)': TFLOPS_PER_TOKEN,
+ 'mfu_6nd_with_attn': new_mfu_6nd_with_attn,
+ 'mfu':new_mfu_6nd_with_attn,
+ 'moe_param_numel': f'{moe_param_numel/1e9:.2f}B',
+ })
all_result['flops_info'] = {
'theory_flops': theory_flops,
+ # 'theory_flops_per_token': theory_flops_per_token,
'model_flops': model_flops,
}
-
- # new_result = self._analysis_single_iter_cost_impl2()
- # all_result['new_duration_time_per_iter'] = new_result['duration_time_per_iter']
- # all_result['sherry'] = new_result
+ all_result['param_numel_info'] = {
+ "dense" : f'{dense_param_numel/1e9:.2f}B',
+ "moe" : f'{moe_param_numel/1e9:.2f}B',
+ "all" : f'{(dense_param_numel+moe_param_numel)/1e9:.2f}B',
+ }
+
+ if self.model_config.model_type == 'moe':
+ activaton_params_numel = dense_param_numel + moe_param_numel * (self.model_config.topk / self.model_config.expert_num)
+ activaton_ratio = activaton_params_numel/(dense_param_numel+moe_param_numel)
+ all_result['param_numel_info'].update({
+ "activations" : f'{activaton_params_numel/1e9:.2f}B',
+ "activations_ratio" : f'{activaton_ratio*100:.2f}%',
+ }
+ )
+
# convert to format
convert_final_result_to_human_format(all_result)
return all_result
@@ -1486,10 +1124,10 @@ def _run(self):
target_point=self.debug_points,
path_list=[],
)
- _ = self.model_chunk_dict["first_stage_chunk"](
+ _ = self.model_chunk_dict[FIRST_CHUNK](
input_info_first_stage, self.path_debug_context
)
- self.pp_state_peak_point["first_stage_chunk"] = self.model_chunk_dict["first_stage_chunk"].compute_activations()
+ self.pp_state_peak_point[FIRST_CHUNK] = self.model_chunk_dict[FIRST_CHUNK].compute_activations()
if self.strategy.pp_size > 2:
seq_len = (
self.strategy.seq_len // self.strategy.tp_size
@@ -1513,10 +1151,10 @@ def _run(self):
# target_point=self.debug_points_last_stage,
path_list=[],
)
- _ = self.model_chunk_dict["middle_stage_chunk"](
+ _ = self.model_chunk_dict[MIDDLE_CHUNK](
input_info_last_stage, self.path_debug_context_last_stage
)
- self.pp_state_peak_point["middle_stage_chunk"] = self.model_chunk_dict["middle_stage_chunk"].compute_activations()
+ self.pp_state_peak_point[MIDDLE_CHUNK] = self.model_chunk_dict[MIDDLE_CHUNK].compute_activations()
if self.strategy.pp_size > 1:
seq_len = (
self.strategy.seq_len // self.strategy.tp_size
@@ -1540,10 +1178,10 @@ def _run(self):
target_point=self.debug_points_last_stage,
path_list=[],
)
- _ = self.model_chunk_dict["last_stage_chunk"](
+ _ = self.model_chunk_dict[LAST_CHUNK](
input_info_last_stage, self.path_debug_context_last_stage
)
- self.pp_state_peak_point["last_stage_chunk"] = self.model_chunk_dict["last_stage_chunk"].compute_activations()
+ self.pp_state_peak_point[LAST_CHUNK] = self.model_chunk_dict[LAST_CHUNK].compute_activations()
def get_pp_stage_peak_mem(self, mem_result, peak_mem_key, toG:bool = False):
assert peak_mem_key in ["peak_mem_with_reserved", "peak_mem"], f"peak_mem_key should be in ['peak_mem_with_reserved', 'peak_mem'] but got {peak_mem_key}"
@@ -1666,7 +1304,7 @@ def search_max_micro_batch_size_fixed_gbs(self, pp_size, dp_size, global_batch_s
if global_batch_size % micro_batch_size != 0 or micro_batch_num < pp_size:
continue
- self.strategy.micro_batch_num = micro_batch_num # TODO(sherry): 固定global_batch_size
+ self.strategy.micro_batch_num = micro_batch_num
self.strategy.micro_batch_size = micro_batch_size
# run
@@ -1718,137 +1356,129 @@ def search_max_micro_batch_size_fixed_gbs(self, pp_size, dp_size, global_batch_s
return all_search_micro_batch_size, all_search_micro_batch_num, all_peak_cached_mem_list, all_cost_list
- def search_selective_recompute(self, global_batch_size, use_reserved_memory):
- # self.search_selective_recompute(global_batch_size=global_batch_size, use_reserved_memory=True, search_best = True, best_mfu=cur_best_mfu, all_search_result=all_search_result)
- """
- Fixed global_batch_size and search for all (selective recompute, micro_batch_size, micro_batch_count) combinations of MFU under the current parallel strategy.
- :param global_batch_size: global batch size
- :param use_reserved_memory: whether to use reserved memory to avoid OOM
- """
+ def log_available_strategy(self, mfu, peak_mem):
+ print(f"Find result parallelism={self.strategy.parallelism}, pp_num_layers={self.get_pp_num_layers()}, recompute={self.strategy.recompute_status},mfu={mfu} gbs={self.strategy.global_batch_size} peak_cached_mem_bytes={peak_mem}GB", flush=True)
- self.strategy.recompute_granularity = "selective_recompute"
- PEAK_MEM_KEY = "peak_mem_with_reserved" if use_reserved_memory else "peak_mem"
- from itertools import product
- all_search_strategy = {}
- params = ['attn_recompute', 'mla_rms_recompute', 'mlp_recompute', 'mlp_rms_recompute']
- combinations = [dict(zip(params, combo)) for combo in product([False, True], repeat=4)]
- for pp_size in [1, 2]:
- for recompute_params in combinations:
- if recompute_params['mla_rms_recompute'] and not recompute_params['attn_recompute']:
- continue
- if recompute_params['mlp_rms_recompute'] and not recompute_params['mlp_recompute']:
- continue
- self.strategy.pp_size = pp_size
- self.strategy.attn_recompute = recompute_params['attn_recompute']
- self.strategy.mla_rms_recompute = recompute_params['mla_rms_recompute']
- self.strategy.mlp_recompute = recompute_params['mlp_recompute']
- self.strategy.mlp_rms_recompute = recompute_params['mlp_rms_recompute']
- rm_tmp()
- all_search_micro_batch_size, all_search_micro_batch_num, all_peak_cached_mem_list, all_cost_list = self.search_max_micro_batch_size_fixed_gbs(pp_size, self.strategy.dp_size, global_batch_size, use_reserved_memory=True, save_all=True)
-
- if len(all_search_micro_batch_size) > 0:
- for search_micro_batch_size, search_micro_batch_num, peak_mem_list, cost_result in zip(all_search_micro_batch_size, all_search_micro_batch_num, all_peak_cached_mem_list, all_cost_list):
- print("find!", peak_mem_list)
- best_strategy = {}
- best_strategy['tp_size'] = self.strategy.tp_size
- best_strategy['ep_size'] = self.strategy.ep_size
- best_strategy['pp_size'] = self.strategy.pp_size
- best_strategy['dp_size'] = self.strategy.dp_size
- best_strategy['micro_batch_size'] = search_micro_batch_size
- best_strategy['micro_batch_num'] = search_micro_batch_num
- best_strategy.update(recompute_params)
- best_strategy['best_mfu'] = cost_result.data['mfu']
- best_strategy['best_TFLOPs'] = cost_result.data['throughput per GPU (TFLOP/s/GPU)']
- best_strategy[PEAK_MEM_KEY] = deepcopy(peak_mem_list)
- best_strategy['bw'] = deepcopy(self.system.real_comm_bw)
- merge_dict(best_strategy, all_search_strategy)
- return all_search_strategy
+ def get_pp_num_layers(self):
+ num_layers_per_pp = math.ceil(self.model_config.layer_num/self.strategy.pp_size)
+ pp_num_layers = f'[{num_layers_per_pp}]x{self.strategy.pp_size-1} + [{self.strategy.num_layers_in_last_pipeline_stage}]' if self.strategy.pp_size > 1 else [self.model_config.layer_num]
+ return pp_num_layers
def dump_paralism_and_recompute_perf(self, mem_result, cost_result):
# from pprint import pprint
# pprint(mem_result.data)
+ dtype = 'fp8' if self.strategy.fp8 else 'bf16'
perf = {
'model_name': self.model_config.model_name,
'system': self.system.sys_name,
- 'parallelism': self.strategy.parallelism,
+ 'parallelism': f'{dtype}.dense{self.model_config.dense_layers}.{self.strategy.parallelism}',
'recompute_status': self.strategy.recompute_status,
- 'TGS_per_gpu' : cost_result.data['throughput_per_accelerator'],
+ 'mfu': cost_result.data["mfu_6nd_with_attn"],
'TFLOPS': cost_result.data['throughput per GPU (TFLOP/s/GPU)'],
- 'mfu': cost_result.data["mfu"],
+ 'TGS_per_gpu' : cost_result.data['throughput_per_accelerator'],
'iter_time': cost_result.data["duration_time_per_iter"],
- 'peak_mem': mem_result.data["peak_mem"] if "peak_mem" in mem_result.data else {s:v['peak_mem'] for s,v in mem_result.data.items()}
+ 'peak_mem': mem_result.data["peak_mem"] if "peak_mem" in mem_result.data else {s:v['peak_mem'] for s,v in mem_result.data.items()},
+ 'peak_mem_with_reserved': mem_result.data["peak_mem_with_reserved"] if "peak_mem_with_reserved" in mem_result.data else {s:v['peak_mem_with_reserved'] for s,v in mem_result.data.items()}
}
return perf
-
- def dump_best_strategy(self, best_strategy):
- # pprint(best_strategy)
- perf = {
- }
+
+ def dump_paralism_and_recompute_bw_perf(self, mem_result, cost_result):
+ perf = self.dump_paralism_and_recompute_perf(mem_result, cost_result)
+ perf['comm_bw_info'] = str(deepcopy(self.system.real_comm_bw))
+ # perf['estimate_details'] = {
+ # 'mem_result': str(mem_result),
+ # 'compute_result': str(cost_result),
+ # 'model_arch':str(self.model_chunk_dict),
+ # 'strategy_config': str(self.strategy),
+ # 'system_config': str(self.system),
+ # 'model_config': str(self.model_config)
+ # }
return perf
- def search_best_selective_recompute(self, use_reserved_memory, gmi_error, best_mfu=None, all_search_result = None):
+ def search_best_selective_recompute(self, use_reserved_memory, gmi_error, best_mfu=None, all_search_result = None, save_path = None):
self.strategy.recompute_granularity = "selective_recompute"
+ accelerator_mem_gbytes = self.system.accelerator.mem_gbs - gmi_error # gmi has 6 GB error
+
PEAK_MEM_KEY = "peak_mem_with_reserved" if use_reserved_memory else "peak_mem"
from itertools import product
best_strategy = {}
params = ['attn_recompute', 'mla_rms_recompute', 'mlp_recompute', 'mlp_rms_recompute']
combinations = [dict(zip(params, combo)) for combo in product([False, True], repeat=4)]
+ combinations = [
+ {
+ 'mla_rms_recompute': True,
+ 'attn_recompute': True,
+ 'mlp_rms_recompute': True,
+ 'mlp_recompute': True,
+ },
+ {
+ 'mla_rms_recompute': True,
+ 'attn_recompute': True,
+ 'mlp_rms_recompute': False,
+ 'mlp_recompute': False,
+ },
+ {
+ 'mla_rms_recompute': False,
+ 'attn_recompute': False,
+ 'mlp_rms_recompute': True,
+ 'mlp_recompute': True,
+ },
+ ]
for recompute_params in combinations:
- if recompute_params['mla_rms_recompute'] and not recompute_params['attn_recompute']:
- continue
- if recompute_params['mlp_rms_recompute'] and not recompute_params['mlp_recompute']:
- continue
- self.strategy.attn_recompute = recompute_params['attn_recompute']
- self.strategy.mla_rms_recompute = recompute_params['mla_rms_recompute']
- self.strategy.mlp_recompute = recompute_params['mlp_recompute']
- self.strategy.mlp_rms_recompute = recompute_params['mlp_rms_recompute']
- self.run_estimate()
- mem_result = self.analysis_mem()
- cost_result = self.analysis_cost()
- peak_mem = self.get_pp_stage_peak_mem(mem_result, PEAK_MEM_KEY, True)
- peak_mem = max(peak_mem.values())
- if peak_mem + gmi_error <= self.system.accelerator.mem_gbs:
- cur_perf = self.dump_paralism_and_recompute_perf(mem_result, cost_result)
- if cur_perf['mfu'] > best_mfu:
- best_mfu = cur_perf['mfu']
- best_strategy = cur_perf
- if all_search_result is not None:
- merge_dict(cur_perf, all_search_result)
+ self.strategy.attn_recompute = recompute_params['attn_recompute']
+ self.strategy.mla_rms_recompute = recompute_params['mla_rms_recompute']
+ self.strategy.mlp_recompute = recompute_params['mlp_recompute']
+ self.strategy.mlp_rms_recompute = recompute_params['mlp_rms_recompute']
+
+ self.run_estimate()
+ mem_result = self.analysis_mem()
+ cost_result = self.analysis_cost()
+ peak_mem_list = self.get_pp_stage_peak_mem(mem_result, PEAK_MEM_KEY, toG=True)
+ peak_cached_mem_gbytes = max(peak_mem_list.values())
+ if peak_cached_mem_gbytes <= accelerator_mem_gbytes:
+ cur_perf = self.dump_paralism_and_recompute_bw_perf(mem_result, cost_result)
+ if cur_perf['mfu'] > best_mfu:
+ best_mfu = cur_perf['mfu']
+ best_strategy = cur_perf
+ self.log_available_strategy(cost_result.data['mfu'], peak_cached_mem_gbytes)
+ if save_path is not None:
+ self._dump_memory_and_cost(mem_result, cost_result, save_path)
+ if all_search_result is not None:
+ merge_dict(cur_perf, all_search_result)
return best_strategy
- def search_best_full_recompute_layer_num(self,
+ def search_best_recompute_layer_num(self,
layer_num,
use_reserved_memory: bool,
gmi_error:int,
best_mfu,
- all_search_result:dict):
+ all_search_result:dict,
+ save_path = None):
"""
Searches for the number of full recompute layers of the highest MFU that can be placed in memory under the current micro_batch_size, micro_batch_count, and parallel policies.
Args:
layer_num (int): layer number
use_reserved_memory (bool): whether to use reserved memory
- gmi_error (int): The error between gmi and the actual allocated storage, the current s5000 is 6GB.
+ gmi_error (int): The error between gmi and the actual allocated storage
best_mfu (float): best mfu
all_search_result (dict): all search result
Returns:
dict: search result
"""
-
- accelerator_mem_gbytes = self.system.accelerator.mem_gbs - gmi_error
- # gmi has 6 GB error
-
- self.strategy.recompute_granularity = "full_block"
+ accelerator_mem_gbytes = self.system.accelerator.mem_gbs - gmi_error # gmi has 6 GB error
PEAK_MEM_KEY = "peak_mem_with_reserved" if use_reserved_memory else "peak_mem"
best_strategy = dict()
- left, right = 0, layer_num // self.strategy.pp_size -1
+ left, right = 0, math.ceil(layer_num/self.strategy.pp_size)
# right = min(right, layer_num-1)
ori_recompute_layer_num = self.strategy.recompute_layer_num
while left <= right:
recompute_layer_num = (left + right) // 2
- assert recompute_layer_num < (layer_num // self.strategy.pp_size), f'recompute_layer_num: {recompute_layer_num}, layer_num: {layer_num}, pp_size: {self.strategy.pp_size}'
+
+ max_recompute_layer_num = math.ceil(layer_num / self.strategy.pp_size)
+ assert recompute_layer_num <= max_recompute_layer_num, f'recompute_layer_num: {recompute_layer_num}, max_recompute_layer_num={max_recompute_layer_num}, layer_num: {layer_num}, pp_size: {self.strategy.pp_size}'
self.strategy.recompute_layer_num = recompute_layer_num
rm_tmp()
@@ -1857,82 +1487,52 @@ def search_best_full_recompute_layer_num(self,
cost_result = self.analysis_cost()
peak_mem_list = self.get_pp_stage_peak_mem(mem_result, PEAK_MEM_KEY, toG=True)
peak_cached_mem_gbytes = max(peak_mem_list.values())
-
if peak_cached_mem_gbytes > accelerator_mem_gbytes:
left = recompute_layer_num + 1
else:
- right = recompute_layer_num -1
-
+ right = recompute_layer_num - 1
# Save best search results
if cost_result.data['mfu'] >= best_mfu:
best_mfu = cost_result.data['mfu']
- best_strategy['model_name'] = self.model_config.model_name
- best_strategy['system'] = self.system.sys_name
- best_strategy["num_gpus_per_node"] = self.system.num_per_node
- best_strategy["memory_capacity"] = f"{self.system.accelerator.mem_gbs} GB"
- best_strategy["intra_comm_bw"] = f"{self.system.networks['high_intra_node'].bandwidth.gbps} GB/s"
- best_strategy["inter_comm_bw"] = f"{self.system.networks['inter_node'].bandwidth.gbps} GB/s"
- best_strategy["world_size"] = self.strategy.world_size
- best_strategy['tp_size'] = self.strategy.tp_size
- best_strategy['ep_size'] = self.strategy.ep_size
- best_strategy['pp_size'] = self.strategy.pp_size
- best_strategy['dp_size'] = self.strategy.dp_size
- best_strategy['edp_size'] = self.strategy.edp_size
- best_strategy['etp_size'] = self.strategy.etp_size
- best_strategy['micro_batch_size'] = self.strategy.micro_batch_size
- best_strategy['micro_batch_num'] = self.strategy.micro_batch_num
- best_strategy['recompute_layer_num'] = recompute_layer_num
-
- best_strategy['best_mfu'] = cost_result.data['mfu']
- best_strategy['best_TFLOPs'] = cost_result.data['throughput per GPU (TFLOP/s/GPU)']
- best_strategy['TGS_per_gpu'] = cost_result.data['throughput_per_accelerator']
- best_strategy[PEAK_MEM_KEY] = str(deepcopy(peak_mem_list))
- best_strategy['comm_bw_info'] = str(deepcopy(self.system.real_comm_bw))
- best_strategy['estimate_details'] = {
- 'mem_result': str(mem_result),
- 'compute_result': str(cost_result),
- 'model_arch':str(self.model_chunk_dict),
- 'strategy_config': str(self.strategy),
- 'system_config': str(self.system),
- 'model_config': str(self.model_config)
- }
- print(f"Find result parallelism={self.strategy.parallelism}, recompute={self.strategy.recompute_status},mfu={cost_result.data['mfu']} gbs={self.strategy.global_batch_size} peak_cached_mem_bytes={peak_cached_mem_gbytes}GB")
+ best_strategy = self.dump_paralism_and_recompute_bw_perf(mem_result, cost_result)
+ self.log_available_strategy(cost_result.data['mfu'], peak_cached_mem_gbytes)
+ if save_path is not None:
+ self._dump_memory_and_cost(mem_result, cost_result, save_path)
if all_search_result is not None:
- """
- cur_serach_result = dict()
- cur_serach_result["model_name"] = model_name
- cur_serach_result["arch"] = arch_name
- cur_serach_result["num_gpus_per_node"] = self.system.num_per_node
- cur_serach_result["memory_capacity"] = f"{self.system.accelerator.mem_gbs} GB"
- cur_serach_result["intra_comm_bw"] = f"{self.system.networks['high_intra_node'].bandwidth.gbps} GB/s"
- cur_serach_result["inter_comm_bw"] = f"{self.system.networks['inter_node'].bandwidth.gbps} GB/s"
- cur_serach_result["world_size"] = self.strategy.world_size
- cur_serach_result["data_type"] = "fp8" if self.strategy.fp8 else "bf16"
- cur_serach_result["tp_size"] = self.strategy.tp_size
- cur_serach_result["ep_size"] = self.strategy.ep_size
- cur_serach_result["pp_size"] = self.strategy.pp_size
- cur_serach_result["dp_size"] = self.strategy.dp_size
- cur_serach_result['edp_size'] = self.strategy.edp_size
- cur_serach_result['etp_size'] = self.strategy.etp_size
- cur_serach_result["micro_batch_size"] = self.strategy.micro_batch_size
- cur_serach_result["micro_batch_num"] = self.strategy.micro_batch_num
- cur_serach_result["gbs"] = self.strategy.dp_size * self.strategy.micro_batch_size * self.strategy.micro_batch_num
- cur_serach_result["recompute_layer_num"] = recompute_layer_num
- cur_serach_result["layer_num"] = layer_num
- cur_serach_result["mfu"] = cost_result.data['mfu']
- cur_serach_result['TFLOPs'] = cost_result.data['throughput per GPU (TFLOP/s/GPU)']
- cur_serach_result['comm_bw_info'] = deepcopy(self.system.real_comm_bw)
- cur_serach_result[PEAK_MEM_KEY] = peak_mem_list
- """
- perf = self.dump_paralism_and_recompute_perf(mem_result, cost_result)
+ perf = self.dump_paralism_and_recompute_bw_perf(mem_result, cost_result)
merge_dict(perf, all_search_result)
self.strategy.recompute_layer_num = ori_recompute_layer_num # recompute_layer_num
return best_strategy
- def search_best_strategy_with_full_recompute(self,
+ def search_best_strategy_no_recompute(self, gmi_error, use_reserved_memory, best_mfu, all_search_result, save_path = None):
+ self.strategy.recompute_granularity = None
+ self.strategy.recompute_layer_num = 0
+ accelerator_mem_gbytes = self.system.accelerator.mem_gbs - gmi_error
+ # gmi has 6 GB error
+ PEAK_MEM_KEY = "peak_mem_with_reserved" if use_reserved_memory else "peak_mem"
+ best_strategy = dict()
+ self.run_estimate()
+ mem_result = self.analysis_mem()
+ cost_result = self.analysis_cost()
+ peak_mem_list = self.get_pp_stage_peak_mem(mem_result, PEAK_MEM_KEY, toG=True)
+ peak_cached_mem_gbytes = max(peak_mem_list.values())
+ if peak_cached_mem_gbytes <= accelerator_mem_gbytes:
+ cur_strategy = self.dump_paralism_and_recompute_bw_perf(mem_result, cost_result)
+ merge_dict(cur_strategy, all_search_result)
+
+ if cost_result.data['mfu'] > best_mfu:
+ best_mfu = cost_result.data['mfu']
+ best_strategy = cur_strategy
+ self.log_available_strategy(cost_result.data['mfu'], peak_cached_mem_gbytes)
+ if save_path is not None:
+ self._dump_memory_and_cost(mem_result, cost_result, save_path)
+
+ return best_strategy
+
+ def search_best_parallel_strategy_with_recompute(self,
world_size:int,
gmi_error:int,
micro_batch_size:int,
@@ -1942,15 +1542,15 @@ def search_best_strategy_with_full_recompute(self,
ep_search_list:List = None,
pp_search_list:List = None,
use_etp:bool = False,
- recompute:str = 'full_block',
- dump_path:str=None,
- dump_best_strategy_details:bool=False):
+ recompute_search_type:str = ['no_recompute', 'full_block', 'selective_recompute'],
+ use_reserved_memory: bool = True,
+ dump_path:str=None):
"""
Searches for the optimal combination of parallel strategies (tp/ep/pp) and full recompute layer configuration that maximizes performance under fixed global batch size constraints.
Args:
world_size (int): world size
- gmi_error (int): The error between gmi and the actual allocated storage, the current s5000 is 6GB.
+ gmi_error (int): The error between gmi and the actual allocated storage
micro_batch_size (int): fixed micro batch size
global_batch_size (int): fixed global batch size
all_search_result (dict): all search result of this model, include (tp, ep, pp, recompute_layer_num) combination.
@@ -1962,6 +1562,9 @@ def search_best_strategy_with_full_recompute(self,
# tp: 1 2 4 8
# ep: 1 2 4 8
# 且 layer num整除
+ if not isinstance(recompute_search_type, list):
+ recompute_search_type = [recompute_search_type]
+
layer_num = self.model_config.layer_num
if tp_search_list is None:
tp_search_list = [1, 2, 4, 8] if self.model_config.model_type == "dense" else [1]
@@ -1971,6 +1574,7 @@ def search_best_strategy_with_full_recompute(self,
pp_search_list = list(range(1, layer_num+1))
global_best_strategy = {}
+ best_strategy_cost_path = f"{dump_path}/best_strategy_costs"
print(f"Start search strategy for world_size={world_size}, model_type={self.model_config.model_type}, model_name={self.model_config.model_name}, system={self.system.sys_name}")
print(f"- tp_search_list={tp_search_list}, ep_search_list={ep_search_list}, pp_search_list={pp_search_list}")
@@ -1978,22 +1582,6 @@ def search_best_strategy_with_full_recompute(self,
print(f"- moe_pad_expert_input_to_capacity={self.model_config.moe_pad_expert_input_to_capacity}")
print(f"- capacity={self.model_config.capacity}")
- def distribute(layer_num, pp_size):
- if pp_size > layer_num:
- return []
- quotient = layer_num // pp_size
- remainder = layer_num % pp_size
- if remainder == 0:
- return [[quotient] * pp_size]
- else:
- first = [quotient + remainder] + [quotient] * (pp_size - 1)
- last = [quotient] * (pp_size - 1) + [quotient + remainder]
- if remainder % 2 == 0:
- split = [quotient + remainder // 2] + [quotient] * (pp_size - 2) + [quotient + remainder // 2]
- else:
- split = [quotient + remainder // 2 + 1] + [quotient] * (pp_size - 2) + [quotient + remainder // 2]
- return [first, last, split]
-
global_best_mfu = -1
for tp_size in tp_search_list:
for ep_size in ep_search_list:
@@ -2003,22 +1591,26 @@ def distribute(layer_num, pp_size):
dp_size = world_size // (pp_size * tp_size)
is_ep_valid = dp_size % ep_size == 0
etp_size = tp_size if use_etp else 1
- is_etp_valid = (world_size %(ep_size* etp_size) == 0) and (etp_size*ep_size < self.system.num_per_node) # TODO(sherry): 临时限定etp_size*ep_size < self.system.num_per_node
+ is_etp_valid = (world_size %(ep_size* etp_size) == 0) and (etp_size*ep_size < self.system.num_per_node)
is_etp_valid = (world_size %(ep_size* etp_size) == 0) # TODO(sherry): 临时限定etp_size*ep_size < self.system.num_per_node
- if is_dp_valid and is_tp_valid and is_ep_valid and is_etp_valid:
- layer_distributes = distribute(layer_num, pp_size)
- for layer_distribute in layer_distributes:
- num_layers_in_first_pipeline_stage = layer_distribute[0] if len(layer_distribute) > 1 else None
- num_layers_in_last_pipeline_stage = layer_distribute[-1] if len(layer_distribute) > 2 else None
- # middle_pp_stage_layer_num = layer_distribute[1:-1] if len(layer_distribute) > 2 else None
+ if pp_size > 1:
+ num_layers_per_pp = math.ceil(layer_num/pp_size)
+ is_pp_valid = num_layers_per_pp > 0
+ num_layers_in_last_pipeline_stage = layer_num - (num_layers_per_pp * (pp_size - 1))
+ is_pp_valid = is_pp_valid and num_layers_in_last_pipeline_stage > 0
+ else:
+ num_layers_in_last_pipeline_stage = None
+ is_pp_valid = True
+
+ if is_dp_valid and is_tp_valid and is_ep_valid and is_etp_valid and is_pp_valid:
# set strategy
self.strategy.world_size = world_size
self.strategy.tp_size = tp_size
self.strategy.ep_size = ep_size
self.strategy.pp_size = pp_size
self.strategy.etp_size = etp_size
- self.strategy.num_layers_in_first_pipeline_stage = num_layers_in_first_pipeline_stage
+ self.strategy.num_layers_in_first_pipeline_stage = None
self.strategy.num_layers_in_last_pipeline_stage = num_layers_in_last_pipeline_stage
@@ -2030,27 +1622,50 @@ def distribute(layer_num, pp_size):
self.strategy.micro_batch_size = micro_batch_size
if micro_batch_size != 0 and search_micro_batch_num != 0:
- if recompute == 'full_block':
- search_best_strategy = self.search_best_full_recompute_layer_num(
- layer_num=self.model_config.layer_num,
- use_reserved_memory = True,
- gmi_error=gmi_error,
- best_mfu=global_best_mfu,
- all_search_result=all_search_result)
- elif recompute == 'selective':
- self.strategy.recompute_layer_num = max(self.model_config.layer_num//pp_size, num_layers_in_first_pipeline_stage if num_layers_in_first_pipeline_stage else 0, num_layers_in_last_pipeline_stage if num_layers_in_last_pipeline_stage else 0)
- search_best_strategy = self.search_best_selective_recompute(
- use_reserved_memory=True,
- gmi_error=gmi_error,
- best_mfu=global_best_mfu,
- all_search_result=all_search_result
- )
- else:
- raise NotImplementedError(f'recompute strategy {recompute} not implemented')
-
- if search_best_strategy and 'mfu' in search_best_strategy:
- global_best_strategy = search_best_strategy
- global_best_mfu = search_best_strategy['mfu']
+ for recompute_type in recompute_search_type:
+ if recompute_type == 'no_recompute':
+ self.strategy.recompute_granularity = None
+ self.strategy.recompute_layer_num = 0
+ self.strategy.recompute_variance = True
+ search_best_strategy = self.search_best_strategy_no_recompute(gmi_error=gmi_error,
+ use_reserved_memory=use_reserved_memory,
+ best_mfu=global_best_mfu,
+ all_search_result=all_search_result)
+ elif recompute_type == 'full_block':
+ self.strategy.recompute_granularity = "full_block"
+ self.strategy.recompute_variance = False # megatron-LM's full recompute does not support variance
+ search_best_strategy = self.search_best_recompute_layer_num(
+ layer_num=self.model_config.layer_num,
+ use_reserved_memory = use_reserved_memory,
+ gmi_error=gmi_error,
+ best_mfu=global_best_mfu,
+ all_search_result=all_search_result,
+ save_path=best_strategy_cost_path)
+ elif recompute_type == 'layer_only':
+ # recompute_granularity is defined by user, we only search the best layer num
+ search_best_strategy = self.search_best_recompute_layer_num(
+ layer_num=self.model_config.layer_num,
+ use_reserved_memory = use_reserved_memory,
+ gmi_error=gmi_error,
+ best_mfu=global_best_mfu,
+ all_search_result=all_search_result,
+ save_path=best_strategy_cost_path)
+ elif recompute_type == 'selective_recompute':
+ self.strategy.recompute_granularity = "selective_recompute"
+ self.strategy.recompute_layer_num = math.ceil(layer_num/pp_size)
+ search_best_strategy = self.search_best_selective_recompute(
+ use_reserved_memory=use_reserved_memory,
+ gmi_error=gmi_error,
+ best_mfu=global_best_mfu,
+ all_search_result=all_search_result,
+ save_path=best_strategy_cost_path
+ )
+ else:
+ raise NotImplementedError(f'recompute strategy {recompute_search_type} not implemented')
+
+ if search_best_strategy and 'mfu' in search_best_strategy:
+ global_best_strategy = search_best_strategy
+ global_best_mfu = search_best_strategy['mfu']
if dump_path is not None and len(global_best_strategy) > 0:
model_name = self.model_config.model_name
@@ -2060,28 +1675,49 @@ def distribute(layer_num, pp_size):
if 'peak_mem' in global_best_strategy and isinstance(global_best_strategy['peak_mem'], dict):
global_best_strategy['peak_mem'] = str(global_best_strategy['peak_mem']) # serialize dict to string to avoid csv dump error
- estimate_details_of_best_strategy = global_best_strategy.pop('estimate_details', None)
best_strategy_df = pd.DataFrame(global_best_strategy, index=[0])
- best_strategy_df.to_csv(f"{dump_path}/{model_name}_{system_name}_seq_len{self.strategy.seq_len}_world_size{self.strategy.world_size}_gbs{self.strategy.global_batch_size}_best_strategy.csv")
+ best_strategy_df.to_csv(f"{dump_path}/{model_name}_{system_name}_seqlen{self.strategy.seq_len}_worldsize{self.strategy.world_size}_gbs{self.strategy.global_batch_size}_best_strategy.csv")
print(best_strategy_df)
if all_search_result is not None:
all_search_result_df = pd.DataFrame(all_search_result)
all_search_result_df = all_search_result_df.sort_values(by ='mfu', ascending=False)
- all_search_result_df.to_csv(f"{dump_path}/{model_name}_{system_name}_seq_len{self.strategy.seq_len}_world_size{self.strategy.world_size}_gbs{self.strategy.global_batch_size}_all_search_strategies.csv")
+ all_search_result_df.to_csv(f"{dump_path}/{model_name}_{system_name}_seqlen{self.strategy.seq_len}_worldsize{self.strategy.world_size}_gbs{self.strategy.global_batch_size}_all_search_strategies.csv")
- if dump_best_strategy_details and estimate_details_of_best_strategy:
- save_path = f'{dump_path}/best_strategy_details'
- os.makedirs(save_path, exist_ok=True)
- for k, v in estimate_details_of_best_strategy.items():
- with open(f"{save_path}/{k}.json", "w") as f:
- f.write(v)
return global_best_strategy
- def analysis(self, save_path=None):
+ def _dump_memory_and_cost(self, mem_result:dict, compute_result:dict, save_path:str):
+ print(f"Saving analysis results to {save_path}")
+ os.makedirs(save_path, exist_ok=True)
+ base_info = {}
+ base_info["arch"] = str(self.model_chunk_dict)
+ base_info["all_param"] = self.model_config.param_numel
+ base_info["act_param"] = self.model_config.activated_param_numel
+ with open(f"{save_path}/model_arch", "w") as f:
+ f.write(base_info["arch"])
+ with open(f"{save_path}/base_info.json", "w") as f:
+ f.write(json.dumps(base_info, indent=2, sort_keys=False, ensure_ascii=False))
+
+ with open(f"{save_path}/mem_result.json", "w") as f:
+ f.write(str(mem_result))
+
+ with open(f"{save_path}/compute_result.json", "w") as f:
+ f.write(str(compute_result))
+
+ with open(f"{save_path}/strategy_config.json", "w") as f:
+ f.write(str(self.strategy))
+
+ with open(f"{save_path}/system_config.json", "w") as f:
+ f.write(str(self.system))
+
+ with open(f"{save_path}/model_config.json", "w") as f:
+ f.write(str(self.model_config))
+
+ def analysis(self, save_path=None, console_log = True):
"""Analyze the performance of the model. Return a dictionary containing the results."""
mem_result = self.analysis_mem()
compute_result = self.analysis_cost()
+
if SIMU_CHECK:
save_path = TMP_PATH
if save_path is not None:
@@ -2114,30 +1750,45 @@ def analysis(self, save_path=None):
# print mfu/tflops/peak_mem
peak_mem = mem_result.data["peak_mem"] if 'peak_mem' in mem_result.data else (({s:r['peak_mem'] for s, r in mem_result.data.items()}))
peak_mem_with_reserved = mem_result.data["peak_mem_with_reserved"] if 'peak_mem_with_reserved' in mem_result.data else (({s:r['peak_mem_with_reserved'] for s, r in mem_result.data.items()}))
- tp = self.strategy.tp_size
- ep = self.strategy.ep_size
- pp = self.strategy.pp_size
- print(f'-------------SIMUMAX SUMMARY \033[33mTP={tp},EP={ep},PP={pp}\033[0m -------------')
- print(f'- parallelism = {self.strategy.parallelism}')
- print(f'- recompute = {self.strategy.recompute_status}')
- print(f"- \033[31mdtype = {'fp8' if self.strategy.fp8 else 'bf16'}, grad_reduce = {'bf16' if self.strategy.grad_reduce_in_bf16 else 'fp32'}\033[0m")
- print(f"- system = {self.system.sys_name}")
- print(f"- model = {self.model_config.model_type}")
- print(f"- \033[32mmfu = {compute_result.data['mfu_6nd_with_attn']:.2f}\033[0m")
- print(f"- \033[32mTFLOPS = {compute_result.data['throughput per GPU (TFLOP/s/GPU)']:.2f} (tflops={compute_result.data['flops_info']['theory_flops']}, duration={compute_result.data['duration_time_per_iter']})\033[0m")
- print(f"- TGS_per_gpu = {compute_result.data['throughput_per_accelerator']}")
- print(f"- \033[31mpeak_alloc_mem = {peak_mem}\033[0m")
- # print(f"- peak_alloc_mem_with_reserved = {peak_mem_with_reserved}")
- # print(f'- net = {self.strategy.net} ')
- # print(f"------------------------------------------")
-
+ if console_log:
+ tp = self.strategy.tp_size
+ ep = self.strategy.ep_size
+ pp = self.strategy.pp_size
+ act_info = f", act={compute_result.data['param_numel_info']['activations']}" if self.model_config.model_type == 'moe' else ''
+ print(f"-------------SIMUMAX SUMMARY \033[33m{self.model_config.model_name}({compute_result.data['param_numel_info']['all']}{act_info}) TP={tp},EP={ep},PP={pp}\033[0m -------------")
+ print(f'- parallelism = layer{self.model_config.layer_num}.dense{self.model_config.dense_layers}.{self.strategy.parallelism}')
+ print(f'- recompute = {self.strategy.recompute_status}')
+ print(f"- \033[31mdtype = {'fp8' if self.strategy.fp8 else 'bf16'}, grad_reduce = {'bf16' if self.strategy.grad_reduce_in_bf16 else 'fp32'}\033[0m")
+ print(f"- system = {self.system.sys_name}")
+ print(f"- model_type = {self.model_config.model_type}")
+ print(f"· \033[32mmfu = {compute_result.data['mfu_6nd_with_attn']:.2f}\033[0m")
+ print(f"· \033[32mTFLOPS = {compute_result.data['throughput per GPU (TFLOP/s/GPU)']:.2f}T (tflops={compute_result.data['flops_info']['theory_flops']}, duration={compute_result.data['duration_time_per_iter']})\033[0m")
+ print(f"· \033[32mTFLOPS_PER_TOKEN = {compute_result.data['throughput per GPU per token (TFLOP/s/GPU/token)']:.2f}T, duration={compute_result.data['duration_time_per_iter']})\033[0m")
+ print(f"· \033[31mpeak_alloc_mem = {peak_mem}\033[0m")
+ print(f"- peak_alloc_mem_with_reserved = {peak_mem_with_reserved}")
+ print(f"- TGS_per_gpu = {compute_result.data['throughput_per_accelerator']}")
+ print(f'- net = {self.strategy.net} ')
+ print(f"------------------------------------------")
+
+
+ # capture graph
+ if ENABLE_SIMU_GRAPH:
+ self.capture(save_path)
+ visualize_with_graphviz(os.path.join(save_path, 'model_graph.json'), output_path=os.path.join(save_path, 'computational_graph'))
return {
+ 'model': self.model_config.model_name,
+ 'model_type': self.model_config.model_type,
+ 'params': compute_result.data['param_numel_info']['all'],
+ 'system': self.system.sys_name,
'peak_mem': peak_mem,
+ 'peak_mem_with_reserved': peak_mem_with_reserved,
+ 'duration_time_per_iter': compute_result.data['duration_time_per_iter'],
'TFLOPS': compute_result.data['throughput per GPU (TFLOP/s/GPU)'],
'TGS_per_gpu': compute_result.data['throughput_per_accelerator'],
'mfu': compute_result.data['mfu_6nd_with_attn'],
'parallelism': self.strategy.parallelism,
'recompute': self.strategy.recompute_status,
- 'dtype': f"{'fp8' if self.strategy.fp8 else 'bf16'},grad_reduce in {'bf16' if self.strategy.grad_reduce_in_bf16 else 'fp32'}"
+ 'dtype': f"{'fp8' if self.strategy.fp8 else 'bf16'},grad_reduce in {'bf16' if self.strategy.grad_reduce_in_bf16 else 'fp32'}",
+ 'net': self.strategy.net,
}
\ No newline at end of file
diff --git a/simumax/core/tensor.py b/simumax/core/tensor.py
index 0429a47..8cbfaf9 100644
--- a/simumax/core/tensor.py
+++ b/simumax/core/tensor.py
@@ -11,17 +11,28 @@
int32 = 4,
int64 = 8
)
-@dataclass
class TensorSize:
+ _id_counter = 0
"""record the shape of the tensor"""
-
- shape: Tuple[int, ...]
- dtype: str = "bf16"
+ def __init__(self, shape: Tuple[int, ...], dtype: str = "bf16", grad_fn=None):
+ # self.shape = shape
+ self.shape = [int(i) for i in shape]
+ self.dtype = dtype
+ self.id = TensorSize._id_counter
+ TensorSize._id_counter += 1
+ self._prev = set()
+ if grad_fn is not None and hasattr(grad_fn, 'inputs'):
+ for i in grad_fn.inputs:
+ self._prev.add(i)
@property
def ndim(self):
return len(self.shape)
+ @property
+ def tensors(self):
+ return [self]
+
def size(self, index: int=None) -> int:
"""
Get the size at the specified index in the tuple.
@@ -65,11 +76,15 @@ def view(self, *args):
def __getitem__(self, index: int) -> int:
return self.shape[index]
- def new(self, dim, new_size):
+ def new_with_dim(self, dim, new_size):
new_shape = list(deepcopy(self.shape))
new_shape[dim] = new_size
return TensorSize(new_shape)
+ def new(self):
+ shape = deepcopy(self.shape)
+ return TensorSize(shape)
+
def unsqeeze(self, dim):
# inplace
self.shape.insert(dim, 1)
@@ -121,6 +136,7 @@ def __str__(self):
# return f"TensorSize(shape={self.shape}, dtype={self.dtype}, mem_size={self.mem_size/1024/1024:.4f} MB)"
return f"TensorSize(shape={self.shape}, dtype={self.dtype})"
+FakeTensor = TensorSize
class Float8Tensor(TensorSize):
def __init__(self, shape):
super().__init__(shape)
diff --git a/simumax/core/transformer/dense_module.py b/simumax/core/transformer/dense_module.py
index 45bd253..10f7095 100644
--- a/simumax/core/transformer/dense_module.py
+++ b/simumax/core/transformer/dense_module.py
@@ -7,10 +7,10 @@
AtomModel, LeafModel, FwdQue,
send_next, send_prev, recv_next, recv_prev,
COM_BUFF)
-from simumax.core.config import ModelConfig, StrategyConfig, SystemConfig, AttentionRecomputeConfig, MLPRecomputeConfig
+from simumax.core.config import ModelConfig, StrategyConfig, SystemConfig, AttentionRecomputeConfig, MLPRecomputeConfig, ENABLE_SIMU_GRAPH
from simumax.core.utils import get_rank_group
import simumax.core.transformer.simu_ops as simu_ops
-from simumax.core.transformer.function import ConcatFunction
+from simumax.core.transformer.function import ConcatFunction,SplitFunction
#region ------------------ Atomic module ------------------
@@ -74,13 +74,12 @@ def prefill(self, args, call_stk='', com_buff=None):
@property
def micro_output_grad_size(self):
# [B, S, H]
- batch_size = self.output_info.tensors[0].size(0)
- seq_len = self.output_info.tensors[0].size(1)
- hidden_size = self.output_info.tensors[0].size(2)
+ batch_size = self.output_info_.tensors[0].size(0)
+ seq_len = self.output_info_.tensors[0].size(1)
+ hidden_size = self.output_info_.tensors[0].size(2)
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
# hidden_size = self.input_info.tensors[0].size(2)
@@ -140,21 +139,22 @@ def _comp_leaf_act_info_impl(self):
def _comp_leaf_model_info_impl(self):
weight_numel = self.vocab_size * self.hidden_size
- self._model_info.weight_bytes = weight_numel * self.element_size
- self._model_info.grad_bytes = (
- self._model_info.weight_bytes
+ self._model_info.weight_numel = weight_numel * self.strategy.tp_size # Statistics the parameters of all tp ranks
+ self._model_info.dense_weight_bytes = weight_numel * self.element_size
+ self._model_info.dense_grad_bytes = (
+ self._model_info.dense_weight_bytes
if not self.strategy.use_fp32_accum_grad
else self.dtype_to_element_size["fp32"] * weight_numel
)
- self._model_info.state_bytes = (
+ self._model_info.dense_state_bytes = (
3 * self.dtype_to_element_size["fp32"] * weight_numel
)
if self.strategy.zero_state >= 1:
- self._model_info.state_bytes /= self.strategy.dp_size
+ self._model_info.dense_state_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 2:
- self._model_info.grad_bytes /= self.strategy.dp_size
+ self._model_info.dense_grad_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 3:
- self._model_info.weight_bytes /= self.strategy.dp_size
+ self._model_info.dense_weight_bytes /= self.strategy.dp_size
def _comp_leaf_flops_info(self):
self._compute_info.fwd_flops = 0
@@ -206,16 +206,21 @@ def __init__(
strategy: StrategyConfig,
system: SystemConfig,
enable_fp8:bool = True,
+ is_last_recompute = False,
+ disable_tensor_parallel = False,
specific_name='ColumnParallelLinear'
) -> None:
super().__init__(input_size, output_size, strategy, system, specific_name)
assert output_size % self.strategy.tp_size == 0
self.layer_idx = layer_idx
self.input_size = input_size
- self.output_size = output_size // self.strategy.tp_size # Split the output size across tp_size, tensor parallel on attention heads dimension
+ self.output_size = output_size if disable_tensor_parallel else output_size // self.strategy.tp_size # Split the output size across tp_size, tensor parallel on attention heads dimension
self.use_bias = use_bias # FIXME(for now unless)
self.has_cached_inputs = has_cached_inputs
self.enable_recompute = enable_recompute
+ self.is_last_recompute = is_last_recompute
+ if self.is_last_recompute and self.enable_recompute:
+ self.set_variance_node(True)
if self.strategy.fp8 and enable_fp8:
self.w_dtype = "fp8"
self.a_dtype = "fp8"
@@ -319,13 +324,12 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- batch_size = self.output_info.tensors[0].size(0)
- seq_len = self.output_info.tensors[0].size(1)
+ batch_size = self.output_info_.tensors[0].size(0)
+ seq_len = self.output_info_.tensors[0].size(1)
# hidden_size = self.output_info.tensors[0].size(2)
return batch_size * seq_len * self.output_size
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
if self.strategy.enable_sequence_parallel:
@@ -435,21 +439,22 @@ def _comp_leaf_act_info_impl(self):
def _comp_leaf_model_info_impl(self):
weight_numel = self.input_size * self.output_size
- self._model_info.weight_bytes = weight_numel * self.w_element_size # fp8_enabled = True, w_element_size=1
- self._model_info.grad_bytes = (
- self._model_info.weight_bytes
+ self._model_info.weight_numel = weight_numel * self.strategy.tp_size # Statistics the parameters of all tp ranks
+ self._model_info.dense_weight_bytes = weight_numel * self.w_element_size # fp8_enabled = True, w_element_size=1
+ self._model_info.dense_grad_bytes = (
+ self._model_info.dense_weight_bytes
if not self.strategy.use_fp32_accum_grad
else self.dtype_to_element_size["fp32"] * weight_numel
)
- self._model_info.state_bytes = (
+ self._model_info.dense_state_bytes = (
3 * self.dtype_to_element_size["fp32"] * weight_numel # w/m/v
)
if self.strategy.zero_state >= 1:
- self._model_info.state_bytes /= self.strategy.dp_size
+ self._model_info.dense_state_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 2:
- self._model_info.grad_bytes /= self.strategy.dp_size
+ self._model_info.dense_grad_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 3:
- self._model_info.weight_bytes /= self.strategy.dp_size
+ self._model_info.dense_weight_bytes /= self.strategy.dp_size
def _comp_leaf_flops_info(self):
base_flops = 2 * self.micro_hidden_state_size * self.output_size
@@ -518,6 +523,7 @@ def __init__(
enable_recompute: bool,
strategy: StrategyConfig,
system: SystemConfig,
+ is_last_recompute = False,
specific_name='RowParallelLinear'
) -> None:
@@ -528,6 +534,9 @@ def __init__(
self.output_size = output_size
self.use_bias = use_bias # FIXME(for now unless)
self.enable_recompute = enable_recompute
+ self.is_last_recompute = is_last_recompute
+ if self.is_last_recompute and self.enable_recompute:
+ self.set_variance_node(True)
self.has_cached_inputs = has_cached_inputs
if self.strategy.fp8:
self.w_dtype = "fp8"
@@ -606,15 +615,14 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- batch_size = self.output_info.tensors[0].size(0)
- seq_len = self.output_info.tensors[0].size(1)
- hidden_size = self.output_info.tensors[0].size(2)
+ batch_size = self.output_info_.tensors[0].size(0)
+ seq_len = self.output_info_.tensors[0].size(1)
+ hidden_size = self.output_info_.tensors[0].size(2)
if self.strategy.enable_sequence_parallel:
seq_len *= self.strategy.tp_size
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
if self.strategy.enable_sequence_parallel:
@@ -707,21 +715,22 @@ def _comp_leaf_act_info_impl(self):
def _comp_leaf_model_info_impl(self):
weight_numel = self.input_size * self.output_size
- self._model_info.weight_bytes = weight_numel * self.w_element_size # fp8_enabled = True, w_element_size=1
- self._model_info.grad_bytes = (
- self._model_info.weight_bytes
+ self._model_info.weight_numel = weight_numel * self.strategy.tp_size # Statistics the parameters of all tp ranks
+ self._model_info.dense_weight_bytes = weight_numel * self.w_element_size # fp8_enabled = True, w_element_size=1
+ self._model_info.dense_grad_bytes = (
+ self._model_info.dense_weight_bytes
if not self.strategy.use_fp32_accum_grad
else self.dtype_to_element_size["fp32"] * weight_numel
)
- self._model_info.state_bytes = (
+ self._model_info.dense_state_bytes = (
3 * self.dtype_to_element_size["fp32"] * weight_numel
)
if self.strategy.zero_state >= 1:
- self._model_info.state_bytes /= self.strategy.dp_size
+ self._model_info.dense_state_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 2:
- self._model_info.grad_bytes /= self.strategy.dp_size
+ self._model_info.dense_grad_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 3:
- self._model_info.weight_bytes /= self.strategy.dp_size
+ self._model_info.dense_weight_bytes /= self.strategy.dp_size
def _comp_leaf_flops_info(self):
base_flops = 2 * self.micro_hidden_state_size * self.output_size
@@ -823,13 +832,12 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- batch_size = self.output_info.tensors[0].size(0)
- seq_len = self.output_info.tensors[0].size(1)
- hidden_size = self.output_info.tensors[0].size(2)
+ batch_size = self.output_info_.tensors[0].size(0)
+ seq_len = self.output_info_.tensors[0].size(1)
+ hidden_size = self.output_info_.tensors[0].size(2)
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.input_info.tensors[0].size(2)
@@ -903,21 +911,22 @@ def _comp_leaf_act_info_impl(self):
def _comp_leaf_model_info_impl(self):
weight_numel = self.norm_size
- self._model_info.weight_bytes = weight_numel * self.element_size
- self._model_info.grad_bytes = (
- self._model_info.weight_bytes
+ self._model_info.weight_numel = weight_numel
+ self._model_info.dense_weight_bytes = weight_numel * self.element_size
+ self._model_info.dense_grad_bytes = (
+ self._model_info.dense_weight_bytes
if not self.strategy.use_fp32_accum_grad
else self.dtype_to_element_size["fp32"] * weight_numel
)
- self._model_info.state_bytes = (
+ self._model_info.dense_state_bytes = (
3 * self.dtype_to_element_size["fp32"] * weight_numel
)
if self.strategy.zero_state >= 1:
- self._model_info.state_bytes /= self.strategy.dp_size
+ self._model_info.dense_state_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 2:
- self._model_info.grad_bytes /= self.strategy.dp_size
+ self._model_info.dense_grad_bytes /= self.strategy.dp_size
if self.strategy.zero_state >= 3:
- self._model_info.weight_bytes /= self.strategy.dp_size
+ self._model_info.dense_weight_bytes /= self.strategy.dp_size
def _comp_leaf_flops_info(self):
# ignore memory bound kernel flops for now
@@ -993,7 +1002,6 @@ def extra_repr(self) -> str:
)
return repr_info
-
class Cat(MetaModule):
"""A simple module just to split, cat, reshape tensor"""
def __init__(
@@ -1005,8 +1013,7 @@ def __init__(
super().__init__(strategy, system)
self.output_dim = output_dim
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
output_info = InputOutputInfo(
@@ -1022,9 +1029,9 @@ def prefill(self, args, call_stk='', com_buff=None):
layer.prefill(args, self.call_stk, com_buff=com_buff)
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_act_info_impl(self):
"""
@@ -1062,8 +1069,6 @@ def _comp_cost_info(self):
bwd_grad_w_op="default",
enable_recompute=self.enable_recompute,
)
-
-
class CoreAttention(MetaModule):
"""Scaled Dot-Product Attention"""
@@ -1079,6 +1084,7 @@ def __init__(
strategy: StrategyConfig,
system: SystemConfig,
specific_name='DotProductAttention',
+ is_last_recompute: bool = False,
) -> None:
super().__init__(strategy, system, specific_name)
self.use_math_sdp = use_math_sdp
@@ -1095,6 +1101,9 @@ def __init__(
self.v_head_dim = head_size
self.has_cached_inputs = has_cached_inputs
self.enable_recompute = enable_recompute
+ self.is_last_recompute = is_last_recompute
+ if self.is_last_recompute and self.enable_recompute:
+ self.set_variance_node(True)
def prefill(self, args, call_stk='', com_buff=None):
self.call_stk = call_stk + self.call_stk
@@ -1119,13 +1128,12 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- batch_size = self.output_info.tensors[0].size(0)
- seq_len = self.output_info.tensors[0].size(1)
- hidden_size = self.output_info.tensors[0].size(2)
+ batch_size = self.output_info_.tensors[0].size(0)
+ seq_len = self.output_info_.tensors[0].size(1)
+ hidden_size = self.output_info_.tensors[0].size(2)
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.head_num * self.head_size
@@ -1142,7 +1150,7 @@ def _pre_op(self):
def get_input_shapes_desc(self, stage):
hidden_states = self.input_info.tensors[0]
batch, seq_len = hidden_states.shape[:2]
- qkv_contiguous = False if 's5000' in self.system.sys_name else True # TODO(sherry): get qkv_contiguous by input stride
+ qkv_contiguous = True # TODO(sherry): get qkv_contiguous by input stride
shape_str = f'batch={int(batch)}, seq_len={int(seq_len)}, head_num={int(self.head_num)}, kv_head_num={int(self.kv_head_num)}, qk_head_dim={int(self.head_size)}, v_head_dim={int(self.v_head_dim)}, qkv_contiguous={qkv_contiguous}'
return shape_str
@@ -1154,8 +1162,6 @@ def _comp_leaf_act_info_impl(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.input_info.tensors[0].size(2)
- # TODO(sherry): 基于cache_info 计算 self._act_info.activation_mem_cache
- # self._act_info.activation_mem_cache = self._comp_input_cache_size()
q_size = batch_size * self.head_num * seq_len * self.head_size
# repeat kv
@@ -1196,18 +1202,15 @@ def _comp_leaf_act_info_impl(self):
# The sdp interface of a certain pytorch will use an extra memory,
# not sure if it is fixed now
bwd_soft_factor += 1
- elif self.use_math_sdp and self.system.accelerator.backend == "musa":
- # torch_musa sdp kernel reuse the output grad memory
- bwd_soft_factor -= 1
self._act_info.bwd_peak_prev_cache_mem = (q_size + k_size) * self.element_size
self._act_info.bwd_peak_mem_no_cache = (
bwd_soft_factor * softmax_size * self.element_size
)
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_flops_info(self):
seq_len = self.input_info.tensors[0].size(1)
@@ -1305,18 +1308,13 @@ def __init__(
strategy: StrategyConfig,
system: SystemConfig,
specific_name='DotProductAttention',
+ is_last_recompute: bool = False,
v_head_dim: int=None,
) -> None:
super().__init__(head_size,head_num,kv_head_num,use_math_sdp,use_flash_sdp,
- has_cached_inputs,enable_recompute,strategy,system,specific_name)
+ has_cached_inputs,enable_recompute,strategy,system,specific_name, is_last_recompute)
self.v_head_dim = v_head_dim
- # def set_input_state_info(self, input_info):
- # super().set_input_state_info(input_info)
- # attention_out = self.output_info
- # if self.use_flash_sdp and self.strategy.fp8: #TODO(sherry): FP8并且使用FA，将attention out cache
- # self.save_for_backward(attention_out)
-
#TODO: memory and net usage need to specify while the implement is different from general core attention
def _comp_leaf_flops_info(self):
@@ -1341,8 +1339,7 @@ def _comp_leaf_flops_info(self):
self._compute_info.bwd_grad_act_flops += base_flops
self._compute_info.bwd_grad_w_flops = 0
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.head_num * self.v_head_dim
@@ -1409,9 +1406,6 @@ def _comp_leaf_act_info_impl(self):
# The sdp interface of a certain pytorch will use an extra memory,
# not sure if it is fixed now
bwd_soft_factor += 1
- elif self.use_math_sdp and self.system.accelerator.backend == "musa":
- # torch_musa sdp kernel reuse the output grad memory
- bwd_soft_factor -= 1
self._act_info.bwd_peak_prev_cache_mem = (q_size + k_size) * self.element_size
self._act_info.bwd_peak_mem_no_cache = (
bwd_soft_factor * softmax_size * self.element_size
@@ -1485,17 +1479,16 @@ def micro_hidden_state_size(self):
hidden_size = self.input_info.tensors[0].size(2)
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
output_info = InputOutputInfo(
- tensors=self.input_info.tensors
+ tensors=[t.new() for t in self.input_info.tensors]
)
return output_info
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_act_info_impl(self):
"""
@@ -1543,11 +1536,13 @@ def __init__(
enable_recompute: bool,
strategy: StrategyConfig,
system: SystemConfig,
+ is_weighted_silu: bool = False,
) -> None:
super().__init__(strategy, system)
self.is_fused = is_fused
self.enable_recompute = enable_recompute
self.has_cached_inputs = has_cached_inputs
+ self.is_weighted_silu = is_weighted_silu
def prefill(self, args, call_stk='', com_buff=None):
self.call_stk = call_stk + self.call_stk
@@ -1566,19 +1561,18 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- input_numel = self.output_info.tensors[0].numel()
+ input_numel = self.output_info_.tensors[0].numel()
return input_numel
- @property
- def output_info(self):
+ def create_output_info(self):
hidden_size = self.input_info.tensors[0].size(-1)
assert hidden_size % 2 == 0, "hidden size should be even"
output_size_list = list(self.input_info.tensors[0].shape[:-1]) + [
hidden_size // 2,
]
output_tensor = [TensorSize(shape=tuple(output_size_list))]
- if len(self.input_info.tensors) > 1:
- output_tensor += self.input_info.tensors[1:]
+ # if len(self.input_info.tensors) > 1:
+ # output_tensor += self.input_info.tensors[1:]
output_info = InputOutputInfo(tensors=output_tensor)
return output_info
@@ -1603,10 +1597,16 @@ def _comp_leaf_act_info_impl(self):
self._act_info.bwd_peak_mem_no_cache = input_size + output_size
self._act_info.bwd_peak_prev_cache_mem = 0
+ if self.is_weighted_silu:
+ probs_mem = self.input_info.tensors[1].numel() * 8
+ # self._act_info.activation_mem_cache += probs_mem # probs are cached in the Permutation
+ self._act_info.fwd_peak_mem_no_cache += probs_mem
+ self._act_info.bwd_peak_mem_no_cache += probs_mem
+
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_flops_info(self):
# ignore memory bound kernel flops for now
@@ -1634,6 +1634,11 @@ def _comp_leaf_mem_accessed_info(self):
)
self._compute_info.bwd_grad_w_accessed_mem = 0
+ if self.is_weighted_silu:
+ probs_mem = self.input_info.tensors[1].numel() * self.dtype_to_element_size[self.strategy.dtype]
+ self._compute_info.fwd_accessed_mem += probs_mem
+ self._compute_info.bwd_grad_act_accessed_mem += probs_mem
+
self._compute_info.recompute_accessed_mem = (
self._compute_info.fwd_accessed_mem if self.enable_recompute else 0
)
@@ -1683,8 +1688,7 @@ def micro_output_grad_size(self):
input_numel = self.input_info.tensors[0].numel()
return input_numel
- @property
- def output_info(self):
+ def create_output_info(self):
output_tensor = [deepcopy(self.input_info.tensors[0])]
if len(self.input_info.tensors) > 1:
output_tensor += self.input_info.tensors[1:]
@@ -1711,9 +1715,9 @@ def _comp_leaf_act_info_impl(self):
self._act_info.bwd_peak_prev_cache_mem = 0
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_flops_info(self):
# ignore memory bound kernel flops for now
@@ -1806,8 +1810,7 @@ def prefill(self, args, call_stk='', com_buff=None):
for layer in self.layers:
layer.prefill(args, self.call_stk, com_buff=com_buff)
- @property
- def output_info(self):
+ def create_output_info(self):
output_info = InputOutputInfo(tensors=[TensorSize(shape=(1,))])
return output_info
@@ -1860,9 +1863,9 @@ def _comp_leaf_act_info_impl(self):
self._act_info_with_recomp = self._act_info
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_flops_info(self):
# ignore memory bound kernel flops for now
@@ -1934,7 +1937,6 @@ def _comp_cost_info(self):
enable_recompute=self.enable_recompute,
)
-
class Float8Quantizer(MetaModule):
def __init__(self, enable_recompute, strategy, system, specific_name='', parent_module=None):
super().__init__(strategy, system, specific_name, parent_module)
@@ -1942,8 +1944,7 @@ def __init__(self, enable_recompute, strategy, system, specific_name='', parent_
self.cache_inputs = False
self.cache_outputs = False
- @property
- def output_info(self):
+ def create_output_info(self):
if isinstance(self.input_info, TensorSize):
output_info = InputOutputInfo([Float8Tensor(deepcopy(self.input_info.shape))])
else:
@@ -1965,6 +1966,8 @@ def _comp_leaf_mem_accessed_info(self):
def extra_repr(self):
return f"enable_recompute={self.enable_recompute}"
+
+
#endregion
#region ----------------- Composite module ----------------
@@ -1978,11 +1981,15 @@ def __init__(self,
enable_recompute,
strategy,
system,
+ is_last_recompute = False,
+ disable_tensor_parallel = False,
specific_name='QuantizedColLinear'):
super().__init__(strategy, system, specific_name, parent_module=None)
assert self.strategy.fp8, 'QuantizedColLinear only support fp8'
- self.quntizer = Float8Quantizer(enable_recompute, strategy, system)
- self.linear = LinearCol(layer_idx, input_size, output_size, use_bias, has_cached_inputs, enable_recompute, strategy, system)
+ self.quntizer = Float8Quantizer(enable_recompute=enable_recompute, strategy=strategy, system=system)
+ self.linear = LinearCol(layer_idx, input_size, output_size, use_bias, has_cached_inputs, enable_recompute, strategy, system,
+ is_last_recompute = is_last_recompute,
+ disable_tensor_parallel = disable_tensor_parallel)
def set_breakpoints(self, status):
self.linear.set_breakpoints(status)
@@ -2001,11 +2008,13 @@ def __init__(self,
enable_recompute,
strategy,
system,
+ is_last_recompute = False,
specific_name='QuantizedColLinear'):
super().__init__(strategy, system, specific_name, parent_module=None)
assert self.strategy.fp8, 'QuantizedColLinear only support fp8'
- self.quntizer = Float8Quantizer(enable_recompute, strategy, system)
- self.linear = LinearRow(layer_idx, input_size, output_size, use_bias, has_cached_inputs, enable_recompute, strategy, system)
+ self.quntizer = Float8Quantizer(enable_recompute=enable_recompute, strategy=strategy, system=system)
+ self.linear = LinearRow(layer_idx, input_size, output_size, use_bias, has_cached_inputs, enable_recompute, strategy, system,
+ is_last_recompute = is_last_recompute)
def set_breakpoints(self, status):
self.linear.set_breakpoints(status)
@@ -2023,7 +2032,7 @@ def __init__(
layer_idx,
config: ModelConfig,
enable_recompute: bool,
- attention_recompute: AttentionRecomputeConfig,
+ attention_recompute_conf: AttentionRecomputeConfig,
strategy: StrategyConfig,
system: SystemConfig,
specific_name='',
@@ -2033,7 +2042,7 @@ def __init__(
self.config = config
self.strategy = strategy
self.system = system
- self.attention_recompute = attention_recompute
+ self.attention_recompute_conf = attention_recompute_conf
self.enable_recompute = enable_recompute # for old version
@@ -2059,7 +2068,7 @@ def __init__(
output_size=output_size,
use_bias=False,
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_recompute,
+ enable_recompute=attention_recompute_conf.q_up_recompute,
strategy=strategy,
system=system
)
@@ -2071,16 +2080,18 @@ def __init__(
use_math_sdp=self.strategy.use_math_sdp,
use_flash_sdp=self.strategy.use_flash_sdp,
has_cached_inputs=False,
- enable_recompute=attention_recompute.attn_recompute,
+ enable_recompute=attention_recompute_conf.core_attn_recompute,
strategy=strategy,
- system=system
+ system=system,
+ is_last_recompute = True
)
+ linear_out_input_size = self.config.head_num * self.config.head_size
self.linear_out = LinearRow_(
layer_idx=layer_idx,
- input_size=self.config.hidden_size,
+ input_size=linear_out_input_size,
output_size=self.config.hidden_size,
has_cached_inputs=False,
- enable_recompute=attention_recompute.out_recompute,
+ enable_recompute=attention_recompute_conf.out_recompute,
use_bias=False,
strategy=strategy,
system=system
@@ -2113,13 +2124,12 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- batch_size = self.output_info.tensors[0].size(0)
- seq_len = self.output_info.tensors[0].size(1)
- hidden_size = self.output_info.tensors[0].size(2)
+ batch_size = self.output_info_.tensors[0].size(0)
+ seq_len = self.output_info_.tensors[0].size(1)
+ hidden_size = self.output_info_.tensors[0].size(2)
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.input_info.tensors[0].size(2)
@@ -2136,7 +2146,7 @@ def __init__(
layer_idx,
config: ModelConfig,
enable_recompute: bool,
- attention_recompute: AttentionRecomputeConfig,
+ attention_recompute_conf: AttentionRecomputeConfig,
strategy: StrategyConfig,
system: SystemConfig,
specific_name='',
@@ -2147,11 +2157,11 @@ def __init__(
self.config = config
self.strategy = strategy
self.system = system
- self.attention_recompute = attention_recompute # for old version
+ self.attention_recompute_conf = attention_recompute_conf # for old version
self.enable_recompute = enable_recompute
query_projection_size = self.config.v_head_dim * self.config.head_num
self.q_head_dim = self.config.qk_head_dim + self.config.qk_pos_emb_head_dim
- self.num_attention_heads_per_partition = self.config.head_num // self.strategy.tp_size #TODO(sherry): 强制走tp还是不强制走tp
+ self.num_attention_heads_per_partition = self.config.head_num // self.strategy.tp_size
# only_enable_sdp = False
if self.strategy.recompute_granularity == "sdp_only":
self.recompute_granularity = "submodule"
@@ -2172,7 +2182,7 @@ def __init__(
output_size=self.config.head_num * self.q_head_dim,
use_bias=False,
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_recompute,
+ enable_recompute=attention_recompute_conf.q_up_recompute,
strategy=strategy,
system=system
)
@@ -2183,7 +2193,9 @@ def __init__(
output_size=self.config.q_lora_rank,
use_bias=False,
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_recompute,
+ enable_recompute=attention_recompute_conf.q_down_recompute,
+ is_last_recompute = True,
+ # disable_tensor_parallel=True,
strategy=strategy,
system=system
)
@@ -2193,7 +2205,7 @@ def __init__(
norm_type="rms_norm",
use_fused_norm=self.strategy.use_fused_norm,
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_norm_recompute,
+ enable_recompute=attention_recompute_conf.q_layernorm_recompute,
strategy=strategy,
system=system
)
@@ -2203,7 +2215,7 @@ def __init__(
output_size=self.config.head_num * self.q_head_dim,
use_bias=False,
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_recompute,
+ enable_recompute=attention_recompute_conf.q_up_recompute,
strategy=strategy,
system=system
)
@@ -2214,22 +2226,18 @@ def __init__(
output_size=self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim,
use_bias=False,
has_cached_inputs=True,
- enable_recompute=attention_recompute.qkv_recompute,
+ enable_recompute=attention_recompute_conf.kv_down_recompute,
+ is_last_recompute = True,
strategy=strategy,
system=system
)
-
- if self.strategy.mla_rms_recompute and self.strategy.recompute_granularity == "selective_recompute":
- if self.config.q_lora_rank is not None:
- self.linear_q_down_proj.set_breakpoints(True)
- self.linear_kv_down_proj.set_breakpoints(True)
-
+
self.kv_layernorm = LayerNorm(
norm_size=self.config.kv_lora_rank,
norm_type="rms_norm",
use_fused_norm=self.strategy.use_fused_norm,
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_norm_recompute,
+ enable_recompute=attention_recompute_conf.kv_layernorm_recompute,
strategy=strategy,
system=system
)
@@ -2240,13 +2248,13 @@ def __init__(
self.config.v_head_dim),
use_bias=False,
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_recompute,
+ enable_recompute=attention_recompute_conf.kv_up_recompute,
strategy=strategy,
system=system,
)
self.rotary_pos_emb = RotaryEmbedding(
has_cached_inputs=False,
- enable_recompute=attention_recompute.qkv_recompute,
+ enable_recompute=attention_recompute_conf.rope_recompute,
# enable_recompute=False,
strategy=strategy,
system=system
@@ -2259,7 +2267,8 @@ def __init__(
use_math_sdp=self.strategy.use_math_sdp,
use_flash_sdp=self.strategy.use_flash_sdp,
has_cached_inputs=False,
- enable_recompute=attention_recompute.attn_recompute,
+ enable_recompute=attention_recompute_conf.core_attn_recompute,
+ is_last_recompute = True,
strategy=strategy,
system=system,
v_head_dim=config.v_head_dim
@@ -2270,16 +2279,23 @@ def __init__(
input_size=query_projection_size,
output_size=self.config.hidden_size,
has_cached_inputs=False,
- enable_recompute=attention_recompute.out_recompute, # TODO: enable_out_recompute
+ enable_recompute=attention_recompute_conf.out_recompute, # TODO: enable_out_recompute
use_bias=False,
strategy=strategy,
system=system
)
- if self.linear_out_proj.enable_recompute and self.strategy.recompute_granularity == "selective_recompute":
- self.linear_out_proj.is_breakpoints = True
-
- is_bf16_and_next_recompute = (not self.strategy.fp8 and self.linear_out_proj.enable_recompute)
- self.core_attention.cache_outputs = self.strategy.use_flash_sdp and (self.strategy.fp8 or is_bf16_and_next_recompute) # FIXME(sherry)：bf16
+ if ENABLE_SIMU_GRAPH:
+ ...
+ else:
+ if self.strategy.mla_rms_recompute and self.strategy.recompute_granularity == "selective_recompute":
+ if self.config.q_lora_rank is not None:
+ self.linear_q_down_proj.set_breakpoints(True)
+ self.linear_kv_down_proj.set_breakpoints(True)
+ if self.linear_out_proj.enable_recompute and self.strategy.recompute_granularity == "selective_recompute":
+ self.linear_out_proj.is_breakpoints = True
+
+ is_bf16_and_next_recompute = (not self.strategy.fp8 and self.linear_out_proj.enable_recompute)
+ self.core_attention.cache_outputs = self.strategy.use_flash_sdp and (self.strategy.fp8 or is_bf16_and_next_recompute) # FIXME(sherry)：bf16
def forward(self, hidden_states:TensorSize, path_debug_context:PathDebugContext):
# ref: Megatron-LM/megatron/core/transformer/multi_latent_attention.py:306
@@ -2314,9 +2330,17 @@ def forward(self, hidden_states:TensorSize, path_debug_context:PathDebugContext)
# kv_combined = gather_from_tensor_model_parallel_region(kv_combined)
# kv_compressed:[s, b, 512], k_pos_emb: [s, b, 64]
- kv_compressed, k_pos_emb = simu_ops.split(
- kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1
- )
+ # kv_compressed, k_pos_emb = simu_ops.split(
+ # kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1
+ # )
+ kv_compressed, k_pos_emb = SplitFunction.apply(parent_model=self,
+ enable_recompute=self.attention_recompute_conf.core_attn_recompute,
+ tensor_size = kv_combined,
+ split_size_or_sections=[self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], split_dim=-1,
+ path_debug_context=path_debug_context,
+ name='kv_combined_Split'
+ )
+
# if self.config.sequence_parallel:
# kv_compressed = scatter_to_sequence_parallel_region(kv_compressed)
@@ -2332,7 +2356,14 @@ def forward(self, hidden_states:TensorSize, path_debug_context:PathDebugContext)
)
# k_no_pe: [s, b, n, 128], value: [s, b, n, 128]
- k_no_pe, value = simu_ops.split(kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1)
+ # k_no_pe, value = simu_ops.split(kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1)
+ k_no_pe, value = SplitFunction.apply(parent_model=self,
+ enable_recompute=self.attention_recompute_conf.core_attn_recompute,
+ tensor_size=kv,
+ split_size_or_sections=[self.config.qk_head_dim, self.config.v_head_dim], split_dim=-1,
+ path_debug_context=path_debug_context,
+ name='KV_Split')
+
# [s, b, 64] -> [s, b, 1, 64]
k_pos_emb:TensorSize = simu_ops.unsqueeze(k_pos_emb, 2)
@@ -2344,7 +2375,10 @@ def forward(self, hidden_states:TensorSize, path_debug_context:PathDebugContext)
# key: [s, b, n, 192]
k_pos_emb = k_pos_emb.expand(-1, -1, self.num_attention_heads_per_partition, -1)
- key = ConcatFunction.apply(self, self.attention_recompute.attn_recompute, [k_no_pe, k_pos_emb], dim=-1, path_debug_context=path_debug_context) # There is storage overhead, cat is defined as Module to manage statistics
+ key = ConcatFunction.apply(self, self.attention_recompute_conf.core_attn_recompute,
+ [k_no_pe, k_pos_emb],
+ dim=-1, path_debug_context=path_debug_context,
+ name='K_pos_emb_Concat') # There is storage overhead, cat is defined as Module to manage statistics
# TODO(sherry): contiguous
# query = query.contiguous()
@@ -2355,7 +2389,13 @@ def forward(self, hidden_states:TensorSize, path_debug_context:PathDebugContext)
query = query.view(s, b, n*d)
key = key.view(s, b, n*d)
value = value.view(s, b, n*d2)
- attn_input = simu_ops.cat([query, key, value], dim = -1) # for simuxmax attention input
+ # attn_input = simu_ops.cat([query, key, value], dim = -1) # for simuxmax attention input
+ attn_input = ConcatFunction.apply(parent_model=self, enable_recompute=self.attention_recompute_conf.core_attn_recompute,
+ tensor_sizes=[query, key, value],
+ dim = -1,
+ path_debug_context=path_debug_context,
+ name='QKV_Concat') # for simuxmax attention input
+
attention_out = self.core_attention(attn_input, path_debug_context) # atomic module
out = self.linear_out_proj(attention_out, path_debug_context)
@@ -2380,13 +2420,12 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- batch_size = self.output_info.tensors[0].size(0)
- seq_len = self.output_info.tensors[0].size(1)
- hidden_size = self.output_info.tensors[0].size(2)
+ batch_size = self.output_info_.tensors[0].size(0)
+ seq_len = self.output_info_.tensors[0].size(1)
+ hidden_size = self.output_info_.tensors[0].size(2)
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.input_info.tensors[0].size(2)
@@ -2401,7 +2440,7 @@ class MLP(MetaModule):
def __init__(self, layer_idx,
config:ModelConfig,
enable_recompute:bool,
- mlp_recompute:MLPRecomputeConfig,
+ mlp_recompute_conf:MLPRecomputeConfig,
strategy:StrategyConfig,
system:SystemConfig,
intermediate_size=None
@@ -2412,7 +2451,7 @@ def __init__(self, layer_idx,
self.strategy = strategy
self.system = system
self.enable_recompute = enable_recompute # for old version
- if not mlp_recompute.linear_recompute:
+ if not mlp_recompute_conf.linear_recompute:
self.recompute_granularity = "submodule"
local_intermediate_size = (intermediate_size if intermediate_size is not None
@@ -2426,14 +2465,14 @@ def __init__(self, layer_idx,
LinearRow_ = QuantizedRowLinear if self.strategy.fp8 else LinearRow
# support selective recompute
-
+ enable_recompute = mlp_recompute_conf.shared_linear_recompute if isinstance(layer_idx, str) and ('shareExpert' in layer_idx) else mlp_recompute_conf.linear_recompute
self.linear_fc1 = LinearCol_(
layer_idx=layer_idx,
input_size=self.config.hidden_size,
output_size=intermediate_size,
use_bias=False,
has_cached_inputs=False,
- enable_recompute=mlp_recompute.linear_recompute,
+ enable_recompute=enable_recompute,
strategy=strategy,
system=system,
)
@@ -2442,7 +2481,8 @@ def __init__(self, layer_idx,
input_size=local_intermediate_size,
output_size=self.config.hidden_size,
has_cached_inputs=False,
- enable_recompute=mlp_recompute.linear_recompute,
+ enable_recompute=enable_recompute,
+ is_last_recompute=True,
use_bias=False,
strategy=strategy,
system=system,
@@ -2451,14 +2491,14 @@ def __init__(self, layer_idx,
self.activation_layer = Swiglu(
is_fused=self.strategy.use_fused_swiglu,
has_cached_inputs=False,
- enable_recompute=mlp_recompute.linear_recompute,
+ enable_recompute=enable_recompute,
strategy=strategy,
system=system,
)
else:
self.activation_layer = Gelu(
has_cached_inputs=False,
- enable_recompute=mlp_recompute.linear_recompute,
+ enable_recompute=enable_recompute,
strategy=strategy,
system=system,
)
diff --git a/simumax/core/transformer/function.py b/simumax/core/transformer/function.py
index 0e11e7e..baf68cc 100644
--- a/simumax/core/transformer/function.py
+++ b/simumax/core/transformer/function.py
@@ -3,42 +3,119 @@
from simumax.core.base_struct import MetaModule, InputOutputInfo, PathDebugContext
# from simumax.core.transformer.dense_module import Qunatizer
+class Function:
+ @staticmethod
+ def apply(cls, *args, **kwargs):
+ raise NotImplementedError
class ConcatModule(MetaModule):
- def __init__(self, dim:int = -1, enable_recompute:bool = False, strategy=None, system=None):
+ def __init__(self, dim:int = -1, enable_recompute:bool = False, strategy=None, system=None, name=None):
super().__init__(strategy, system)
self.dim = dim
self.enable_recompute = enable_recompute
self.is_leaf_module = True
-
- @property
- def output_info(self):
+ self.name = name if name else 'ConcatModule'
+
+ def create_output_info(self):
# return TensorSize or InputOutputInfo
tensor_sizes = self.input_info.tensors
if len(tensor_sizes) == 0:
return InputOutputInfo([])
concat_size = sum([t[self.dim] for t in tensor_sizes])
- return tensor_sizes[0].new(self.dim, concat_size)
+ output_info = tensor_sizes[0].new_with_dim(self.dim, concat_size)
+ return output_info
def extra_repr(self) -> str:
repr_info = f"concat_dim={self.dim}, enable_recompute={self.enable_recompute}"
return repr_info
+
+class SplitModule(MetaModule):
+ def __init__(self, split_size_or_sections:List[int], split_dim, enable_recompute:bool = False, strategy=None, system=None, name=None):
+ super().__init__(strategy, system)
+ self.split_size_or_sections = split_size_or_sections
+ self.split_dim = split_dim
+ self.enable_recompute = enable_recompute
+ self.name = name if name else 'SplitModule'
+
+ def create_output_info(self):
+ tensor_size = self.input_info.tensors[0] if isinstance(self.input_info, InputOutputInfo) else self.input_info
+ split_dim = self.split_dim
+ split_size_or_sections = self.split_size_or_sections
+ if isinstance(split_size_or_sections, int):
+ assert tensor_size[split_dim] % split_size_or_sections == 0, f"tensor_size[dim]={tensor_size[split_dim]} split_size_or_sections={split_size_or_sections}"
+ print(f"split_size_or_sections is int, tensor_size[dim] is {tensor_size[split_dim]}, split_size_or_sections is {split_size_or_sections}")
+ split_size_or_sections = [tensor_size[split_dim] // split_size_or_sections] * split_size_or_sections
+
+ assert tensor_size[split_dim] == sum(split_size_or_sections), f"tensor_size[dim]={tensor_size[split_dim]} sum(split_size_or_sections)={sum(split_size_or_sections)}, tensor_size={tensor_size.shape}, split_dim={split_dim}"
+ output_info = InputOutputInfo(tensors=[tensor_size.new_with_dim(split_dim, size) for size in split_size_or_sections])
+ return output_info
- # TODO(sherry): concat的统计函数。。。
+ def extra_repr(self) -> str:
+ repr_info = f"split_dim={self.split_dim}, enable_recompute={self.enable_recompute}"
+ return repr_info
-class Function:
- @staticmethod
- def apply(cls, *args, **kwargs):
- raise NotImplementedError
+class AddModule(MetaModule):
+ def __init__(self, enable_recompute:bool = False, strategy=None, system=None, name=None):
+ super().__init__(strategy, system)
+ self.enable_recompute = enable_recompute
+ self.name = name
+
+ def create_output_info(self):
+ # recover the original input
+ assert self.output_info_ is None
+ output_info = InputOutputInfo(tensors=[self.input_info.tensors[0].new()])
+ return output_info
+
+ def extra_repr(self) -> str:
+ repr_info = f"enable_recompute={self.enable_recompute}"
+ return repr_info
+
+class UnsqueezeModule(MetaModule):
+ def __init__(self, unsqueeze_dim:int, enable_recompute:bool = False, strategy=None, system=None, name=None):
+ super().__init__(strategy, system)
+ self.unsqueeze_dim = unsqueeze_dim
+ self.enable_recompute = enable_recompute
+ self.name = name if name else 'UnsqueezeModule'
+ def create_output_info(self):
+ inputs = self.input_info.tensors[0] if isinstance(self.input_info, InputOutputInfo) else self.input_info
+ outputs = inputs.new()
+ outputs.squeeze(self.unsqueeze_dim)
+ return InputOutputInfo(tensors=outputs)
+
class ConcatFunction(Function):
@staticmethod
- def apply(model:MetaModule, enable_recompute:bool, tensor_sizes: List[TensorSize], dim:int = -1, path_debug_context: PathDebugContext = None):
+ def apply(parent_model:MetaModule, enable_recompute:bool, tensor_sizes: List[TensorSize], dim:int = -1, path_debug_context: PathDebugContext = None, name = None):
# model.output_size = TensorSize.concat(tensor_sizes, dim)
- concat_module = ConcatModule(dim, enable_recompute, model.strategy, model.system)
- concat_module.parent_module = model # Bind parent module
+ concat_module = ConcatModule(dim, enable_recompute, parent_model.strategy, parent_model.system, name)
+ concat_module.parent_module = parent_model # Bind parent module
input_info = InputOutputInfo(tensor_sizes)
out = concat_module(input_info, path_debug_context = path_debug_context) # Reuse the __call__ method of MetaModule, call related functions for statistics, and register the concat_module into parent_module
+ return out
+
+class SplitFunction(Function):
+ @staticmethod
+ def apply(parent_model:MetaModule, enable_recompute:bool, tensor_size:TensorSize, split_size_or_sections:int, split_dim:int = -1, path_debug_context: PathDebugContext = None, name = None):
+ # model.output_size = TensorSize.split(tensor_size, split_size_or_sections, dim)
+ split_module = SplitModule(split_size_or_sections, split_dim, enable_recompute, parent_model.strategy, parent_model.system, name)
+ split_module.parent_module = parent_model # Bind parent module
+
+ input_info = InputOutputInfo([tensor_size]) if isinstance(tensor_size, TensorSize) else tensor_size
+ out = split_module(input_info, path_debug_context = path_debug_context) # Reuse the __call__ method of MetaModule, call related functions for statistics, and register the split_module into parent_module
return out
+
+class AddFunction(Function):
+ @staticmethod
+ def apply(parent_model:MetaModule, enable_recompute:bool, tensor_size1:TensorSize, tensor_size2:TensorSize, path_debug_context: PathDebugContext = None, name = None):
+ # model.output_size = TensorSize.split(tensor_size, split_size_or_sections, dim)
+ add_module = AddModule(enable_recompute, parent_model.strategy, parent_model.system, name)
+ add_module.parent_module = parent_model # Bind parent module
+ if isinstance(tensor_size1, InputOutputInfo):
+ tensor_size1 = tensor_size1.tensors[0]
+ if isinstance(tensor_size2, InputOutputInfo):
+ tensor_size2 = tensor_size2.tensors[0]
+ input_info = InputOutputInfo([tensor_size1, tensor_size2])
+ out = add_module(input_info, path_debug_context = path_debug_context) # Reuse the __call__ method of MetaModule, call related functions for statistics, and register the split_module into parent_module
+ return out
\ No newline at end of file
diff --git a/simumax/core/transformer/language_model.py b/simumax/core/transformer/language_model.py
index bb4ebe2..ae2439b 100644
--- a/simumax/core/transformer/language_model.py
+++ b/simumax/core/transformer/language_model.py
@@ -3,8 +3,8 @@
from copy import deepcopy
from dataclasses import dataclass, asdict
from typing import List
-from simumax.core.base_struct import MetaModule, InputOutputInfo, PathDebugContext, LinearBase, RecomputeStatus
-from simumax.core.config import ModelConfig, StrategyConfig, SystemConfig, AttentionRecomputeConfig, MLPRecomputeConfig, SIMU_DEBUG
+from simumax.core.base_struct import MetaModule, InputOutputInfo, PathDebugContext, LinearBase, RecomputeStatus, RecomputeBreakModule
+from simumax.core.config import ModelConfig, StrategyConfig, SystemConfig, AttentionRecomputeConfig, MLPRecomputeConfig, SIMU_DEBUG, ENABLE_SIMU_GRAPH
from simumax.core.transformer.dense_module import Embedding, Attention, MLAAttention, LayerNorm, LinearCol, MLP, ParallelCE
from simumax.core.transformer.moe_module import ExpertMLP
@@ -41,7 +41,7 @@ def set_peak(self, path, mem, stage:str):
self.recomp_bwd_peak_mem = mem
def update_peak(self, path, mem, stage:str):
- if mem > self.peak_mem:
+ if mem >= self.peak_mem:
self.set_peak(path, mem, stage)
def set_stage(self, stage):
@@ -126,7 +126,7 @@ def __init__(
norm_type="rms_norm",
use_fused_norm=self.strategy.use_fused_norm,
has_cached_inputs=False,
- enable_recompute=attention_recompute.input_norm_recompute,
+ enable_recompute=attention_recompute.input_layernorm_recompute,
strategy=strategy,
system=system,
)
@@ -140,7 +140,7 @@ def __init__(
layer_idx=layer_idx,
config=self.config,
enable_recompute=enable_attn_recompute, # for old version
- attention_recompute = attention_recompute,
+ attention_recompute_conf = attention_recompute,
strategy=strategy,
system=system,
specific_name='SelfAttention'
@@ -150,7 +150,7 @@ def __init__(
layer_idx=layer_idx,
config=self.config,
enable_recompute=enable_attn_recompute, # for old version
- attention_recompute = attention_recompute,
+ attention_recompute_conf = attention_recompute,
strategy=strategy,
system=system,
specific_name='SelfAttention'
@@ -173,7 +173,7 @@ def __init__(
layer_idx=layer_idx,
config=self.config,
enable_recompute=enable_mlp_recompute, # for old version
- mlp_recompute = mlp_recompute,
+ mlp_recompute_conf = mlp_recompute,
strategy=strategy,
system=system,
)
@@ -187,7 +187,6 @@ def __init__(
system=system,
specific_name='MoELayer'
)
- # if self.layernorm_input.enable_recompute:
def forward(self, input_info:InputOutputInfo, path_debug_context:PathDebugContext):
hidden_state = self.layernorm_input(input_info, path_debug_context)
hidden_state = self.attention(hidden_state, path_debug_context)
@@ -234,9 +233,9 @@ def __init__(
)
# self.layers = []
for i in range(layer_num):
- enable_recompute = self.strategy.enable_recompute and (
+ enable_recompute = self.strategy.is_recompute and (
i < self.strategy.recompute_layer_num
- )
+ )
attention_recompute_config = self.strategy.parse_attention_recompute(i)
mlp_recompute_config = self.strategy.parse_mlp_recompute(i)
setattr(self, f'layer_{i}', LLMBlock(
@@ -290,6 +289,7 @@ def add_ordered_module_hook(p_module:MetaModule, sub_module:MetaModule):
cur_m = sub_module
if cur_m.is_leaf_module:
+ cur_m.call_idx = len(self.all_leaf_nodes)
self.all_leaf_nodes.append(cur_m)
# set default recompute status
@@ -359,7 +359,7 @@ def _comp_fwd_activations(self, enable_recompute, compute_nodes:List[MetaModule
act_info.cache_for_bwd_mem = act_info.total_activation_mem_cache
global_cache_mem += act_info.cache_for_bwd_mem
elif stage == "forward" and m.recompute_status == RecomputeStatus.FIRST:
- act_info.cache_for_bwd_mem = m.all_input_element_num()
+ act_info.cache_for_bwd_mem = m.all_input_element_num() if not m.offload_inputs else 0
global_cache_mem += act_info.cache_for_bwd_mem
else:
@@ -368,7 +368,7 @@ def _comp_fwd_activations(self, enable_recompute, compute_nodes:List[MetaModule
if stage == "forward" and enable_recompute:
# if m.recompute_status in [RecomputeStatus.FIRST, RecomputeStatus.LAST, RecomputeStatus.MIDDLE, RecomputeStatus.NO_RECOMPUTE]:
- if m.recompute_status in [RecomputeStatus.FIRST, RecomputeStatus.MIDDLE, RecomputeStatus.LAST]:
+ if m.recompute_status in [RecomputeStatus.FIRST, RecomputeStatus.LAST]:
if SIMU_DEBUG:
print(f"Find {m.recompute_status} node: {m.full_name}")
peak_point.update_peak(f"{m.full_name}: {m.current_full_module_path}", global_cache_mem, stage)
@@ -422,7 +422,8 @@ def _comp_bwd_activations(self, enable_recompute, global_cache_mem:int = 0, peak
prepare_recompute_ready = False
else:
act_info = m.get_act_info()
- cur_peak_mem = global_cache_mem + m.all_input_element_num() if is_last_module else act_info.bwd_peak_mem_no_cache
+ cur_peak_mem = global_cache_mem + (m.all_input_element_num() if is_last_module else act_info.bwd_peak_mem_no_cache)
+ # cur_peak_mem = global_cache_mem + m.all_input_element_num() if is_last_module else act_info.bwd_peak_mem_no_cache
peak_point.update_peak(f"{m.full_name}: {m.current_full_module_path}", cur_peak_mem, "backward")
@@ -443,39 +444,23 @@ def _comp_bwd_activations(self, enable_recompute, global_cache_mem:int = 0, peak
def compute_activations(self):
leaf_nodes = self.get_all_leaf_modules()
self.set_breakpoints(leaf_nodes)
-
- no_recompute_peak_point = PeakPoint()
- with_recompute_peak_point = PeakPoint()
-
- # no recompute
- enable_recompute = False
+ peak_point = PeakPoint()
+ enable_recompute = self.strategy.enable_recompute
global_cache_mem = self._comp_fwd_activations(enable_recompute=enable_recompute,
compute_nodes=leaf_nodes,
global_cache_mem=0,
- peak_point=no_recompute_peak_point)
+ peak_point=peak_point)
global_cache_mem = self._comp_bwd_activations(enable_recompute=enable_recompute,
global_cache_mem=global_cache_mem,
- peak_point=no_recompute_peak_point)
+ peak_point=peak_point)
for i, m in enumerate(leaf_nodes):
assert m._act_info.cache_for_bwd_mem == 0, f"{leaf_nodes[i].full_name}._act_info.cache_for_bwd_mem should be 0, but is {m._act_info.cache_for_bwd_mem/1024/1024:.2f} MB"
assert global_cache_mem == 0, f"global_cache_mem should be 0, but is {global_cache_mem/1024/1024:.2f} MB"
- # with recompute
- enable_recompute = True
- with_recompute_global_cache_mem = self._comp_fwd_activations(enable_recompute=enable_recompute,
- compute_nodes = leaf_nodes,
- global_cache_mem=0,
- peak_point=with_recompute_peak_point)
- with_recompute_global_cache_mem = self._comp_bwd_activations(enable_recompute=enable_recompute,
- global_cache_mem=with_recompute_global_cache_mem,
- peak_point=with_recompute_peak_point)
- assert with_recompute_global_cache_mem == 0, f"with_recompute_global_cache_mem should be 0, but is {with_recompute_global_cache_mem/1024/1024:.2f} MB"
-
- return dict(no_recompute_peak_point=no_recompute_peak_point,
- with_recompute_peak_point=with_recompute_peak_point)
+ return peak_point
def get_all_gemm_cost_info(self):
all_gemm_info = {
@@ -524,7 +509,7 @@ def get_all_gemm_cost_info(self):
def analysis_op_info(self, return_details=False):
assert self.init_ready and self.input_info and self.status_ready, "Please initialize the model first!"
- leaf_moduls = self.get_all_leaf_modules()
+ leaf_modules = self.get_all_leaf_modules()
op_infos = {
"op":[],
"input_shapes":[],
@@ -540,9 +525,9 @@ def analysis_op_info(self, return_details=False):
op_infos['compute_only_details'] = []
op_infos['IO_details'] = []
- for m in leaf_moduls:
+ for m in leaf_modules:
# forward
- output_shape = m.output_info.shapes if isinstance(m.output_info, InputOutputInfo) else [m.output_info.shape]
+ output_shape = m.output_info_.shapes if isinstance(m.output_info_, InputOutputInfo) else [m.output_info_.shape]
op_infos["op"].append(m.__class__.__name__)
op_infos["input_shapes"].append(m.input_info.shapes + ([m.weight.shape] if hasattr(m, "weight") else []))
op_infos['output_shapes'].append(output_shape)
@@ -590,7 +575,7 @@ def analysis_op_info(self, return_details=False):
else:
d_w_lhs_shape = [m.input_info.shapes]
op_infos["op"].append(m.__class__.__name__ + "_bwd_w")
- op_infos["input_shapes"].append(d_w_lhs_shape + m.output_info.shapes)
+ op_infos["input_shapes"].append(d_w_lhs_shape + m.output_info_.shapes)
op_infos['output_shapes'].append([m.get_weight().shape])
op_infos["flops"].append(m._compute_info.bwd_grad_w_flops)
op_infos["IO"].append(m._compute_info.bwd_grad_w_accessed_mem)
diff --git a/simumax/core/transformer/moe_module.py b/simumax/core/transformer/moe_module.py
index dcfccc4..5fb0c82 100644
--- a/simumax/core/transformer/moe_module.py
+++ b/simumax/core/transformer/moe_module.py
@@ -13,7 +13,8 @@
import simumax.core.transformer.simu_ops as simu_ops
from simumax.core.transformer.dense_module import Swiglu, Gelu, MLP, Float8Quantizer, LinearCol, LinearRow
from simumax.core.utils import get_rank_group
-
+from simumax.core.transformer.function import AddFunction
+Input = InputOutputInfo
#region ------------------ Atomic module ------------------
class Router(MetaModule):
"""
@@ -89,15 +90,13 @@ def local_logits_size(self):
ep_num = self.expert_num
return b * seq_len * ep_num
- @property
- def output_info(self):
+ def create_output_info(self):
# FIXME(sherry): check this, return [hidden_states, scores, routting_map]
batch_size = self.input_info.tensors[0].size(0)
seq_len = self.input_info.tensors[0].size(1)
- hidden_size = self.input_info.tensors[0].size(2)
+ # hidden_size = self.input_info.tensors[0].size(2)
hidden_states = InputOutputInfo(
- tensors=[TensorSize(shape=(batch_size, seq_len, hidden_size)),
- TensorSize(shape=(batch_size, seq_len, self.expert_num), dtype="int32")]
+ tensors=[TensorSize(shape=(batch_size, seq_len, self.expert_num), dtype="int32")]
)
return hidden_states
@@ -144,7 +143,7 @@ def _comp_leaf_act_info_impl(self):
seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.input_info.tensors[0].size(2)
input_size = batch_size * seq_len * hidden_size
- self._act_info.activation_mem_cache = self.input_info.tensors[0].numel() * self.element_size
+ self._act_info.activation_mem_cache = input_size * self.element_size
if self.has_cached_inputs:
self._act_info.activation_mem_cache = 0
@@ -167,21 +166,22 @@ def _comp_leaf_model_info_impl(self):
weight = input(linear), scores([S,K], softmax)
"""
weight_numel = self.hidden_size * self.expert_num
- self._model_info.weight_bytes = weight_numel * self.element_size
- self._model_info.grad_bytes = (
- self._model_info.weight_bytes
+ self._model_info.weight_numel = weight_numel
+ self._model_info.dense_weight_bytes = weight_numel * self.element_size
+ self._model_info.dense_grad_bytes = (
+ self._model_info.dense_weight_bytes
if not self.strategy.use_fp32_accum_grad
else self.dtype_to_element_size["fp32"] * weight_numel
)
- self._model_info.state_bytes = (
+ self._model_info.dense_state_bytes = (
3 * self.dtype_to_element_size["fp32"] * weight_numel
)
if self.strategy.zero_state >= 1:
- self._model_info.state_bytes /= self.strategy.edp_size
+ self._model_info.dense_state_bytes /= self.strategy.edp_size
if self.strategy.zero_state >= 2:
- self._model_info.grad_bytes /= self.strategy.edp_size
+ self._model_info.dense_grad_bytes /= self.strategy.edp_size
if self.strategy.zero_state >= 3:
- self._model_info.weight_bytes /= self.strategy.edp_size
+ self._model_info.dense_weight_bytes /= self.strategy.edp_size
def _comp_leaf_flops_info(self):
# Count Gating
@@ -387,8 +387,14 @@ def input_act_size(self):
hidden_size = self.input_info.tensors[0].size(2)
return batch_size * seq_len * hidden_size
- @property
- def output_info(self):
+ def _pre_op(self):
+ super()._pre_op()
+ # if self.strategy.dispatch_probs:
+ # assert len(self.input_info.tensors) == 2, "dispatch_probs=True requires two inputs in Permutation, [x, probs]"
+ # else:
+ # assert len(self.input_info.tensors) == 1, "dispatch_probs=False requires one inputs in Permutation, [x]"
+
+ def create_output_info(self):
batch_size = self.input_info.tensors[0].size(0)
part_seq_len = self.input_info.tensors[0].size(1)
hidden_size = self.input_info.tensors[0].size(2)
@@ -406,9 +412,6 @@ def output_info(self):
TensorSize(
shape=(balance_token_num, hidden_size)
), # permuted moe input
- # TensorSize(
- # shape=(batch_size, part_seq_len, hidden_size)
- # ), # original moe input
]
)
return output_info
@@ -424,7 +427,7 @@ def _comp_leaf_intra_net_info(self):
comm_size,
comm_num=self.strategy.tp_size,
net=self.strategy.tp_net,
- comm_stage="Permutation_FWD_EP"
+ comm_stage="Dispatch_FWD_EP"
)
# bwd
self._cost_info.bwd_grad_act_net_time += self.system.compute_net_op_time(
@@ -432,7 +435,7 @@ def _comp_leaf_intra_net_info(self):
comm_size,
comm_num=self.strategy.tp_size,
net=self.strategy.tp_net,
- comm_stage="Permutation_BWD_EP"
+ comm_stage="Dispatch_BWD_EP"
)
if self.strategy.ep_size > 1:
comm_size = (
@@ -444,16 +447,37 @@ def _comp_leaf_intra_net_info(self):
comm_size,
comm_num=self.strategy.ep_size,
net=self.strategy.ep_net,
- comm_stage="Permutation_FWD_EP"
+ comm_stage="Dispatch_FWD_EP"
)
+
# bwd
self._cost_info.bwd_grad_act_net_time += self.system.compute_net_op_time(
"all2all",
comm_size,
comm_num=self.strategy.ep_size,
net=self.strategy.ep_net,
- comm_stage="Permutation_BWD_EP"
+ comm_stage="Dispatch_BWD_EP"
)
+
+ # HACK(sherry): all2all the router probs to expert, and fused combined probs to SiluOp in ExpertMLP, to avoid the activation_mem_cache in Unpermutaion
+ if self.strategy.dispatch_probs:
+ prob_comm_size = self.input_info.tensors[1].numel() * self.dtype_to_element_size[self.strategy.dtype]
+ self._cost_info.fwd_net_time += self.system.compute_net_op_time(
+ "all2all",
+ prob_comm_size,
+ comm_num=self.strategy.ep_size,
+ net=self.strategy.ep_net,
+ comm_stage="Permutation_FWD_EP_PROB"
+ )
+ self._cost_info.bwd_grad_act_net_time += self.system.compute_net_op_time(
+ "all2all",
+ prob_comm_size,
+ comm_num=self.strategy.ep_size,
+ net=self.strategy.ep_net,
+ comm_stage="Permutation_BWD_EP_PROB"
+ )
+ # HACK(sherry)
+
if self.strategy.etp_size > 1:
comm_size = (
self.permuted_act_size
@@ -480,16 +504,17 @@ def _comp_leaf_intra_net_info(self):
self._cost_info.recompute_net_time = self._cost_info.fwd_net_time
def _comp_leaf_act_info_impl(self):
- self._act_info.activation_mem_cache = self.input_info.tensors[1].numel() * 8
+ probs_mem = self.input_info.tensors[1].numel() * 8
+ self._act_info.activation_mem_cache = probs_mem
self._act_info.fwd_peak_mem_no_cache = 0
self._act_info.fwd_peak_prev_cache_mem = 0
self._act_info.bwd_peak_mem_no_cache = 0
self._act_info.bwd_peak_prev_cache_mem = 0
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_flops_info(self):
"""
@@ -569,6 +594,7 @@ def __init__(
self.has_cached_inputs = has_cached_inputs
self.enable_recompute = enable_recompute
self.moe_dispatcher_policy = moe_dispatcher_policy
+ self.ori_shape = None
def prefill(self, args, call_stk='', com_buff=None):
self.call_stk = call_stk + self.call_stk
@@ -688,13 +714,24 @@ def act_size_before_combined(self):
@property
def act_size_after_combined(self):
# only consider balanced case
- act_size = self.output_info.tensors[0].numel()
+ act_size = self.output_info_.tensors[0].numel()
return act_size
+
+ def _pre_op(self):
+ super()._pre_op()
+ if not self.strategy.dispatch_probs:
+ assert len(self.input_info.tensors) == 2, "dispatch_probs=False requires two inputs in Permutation, [x, probs]"
+ else:
+ assert len(self.input_info.tensors) == 1, "dispatch_probs=True requires one inputs in Permutation, [x]"
+ def set_ori_shape(self, shape):
+ self.ori_shape = shape
- @property
- def output_info(self):
+ def create_output_info(self):
# recover the original input
- output_info = InputOutputInfo(tensors=[self.input_info.tensors[1]])
+ assert self.output_info_ is None
+ assert self.ori_shape is not None
+ output_info = InputOutputInfo(tensors=[TensorSize(shape=self.ori_shape)])
+ # print('-- unpermute output_info', output_info)
return output_info
def _comp_leaf_intra_net_info(self):
@@ -771,16 +808,26 @@ def _comp_leaf_act_info_impl(self):
"""
Mainly layout operators, ignore for now
"""
- self._act_info.activation_mem_cache = self.input_info.tensors[0].numel() * self.element_size
- self._act_info.fwd_peak_mem_no_cache = self.input_info.tensors[0].numel() * self.element_size
+ # HACK(sherry): the weighted_probs is fused in SiluOP.
+ # if dispatch_probs=True, no cache.
+ # if dispatch_probs=False, cache the unpermute_before_hidden_states and probs (for mul op).
+ if self.strategy.dispatch_probs:
+ self._act_info.activation_mem_cache = 0
+ self._act_info.fwd_peak_mem_no_cache = max(self.act_size_before_combined, self.act_size_after_combined) * self.element_size
+ self._act_info.bwd_peak_mem_no_cache = 0
+ else:
+ # mul
+ self._act_info.activation_mem_cache = self.act_size_before_combined * self.element_size # Cache hidden states, probs are cache in Permutation
+ self._act_info.fwd_peak_mem_no_cache = self.act_size_before_combined * self.element_size + self.act_size_after_combined * self.element_size
+ self._act_info.bwd_peak_mem_no_cache = self.act_size_before_combined * self.element_size + self.act_size_after_combined * self.element_size
+ # HACK(sherry)
self._act_info.fwd_peak_prev_cache_mem = 0
- self._act_info.bwd_peak_mem_no_cache = 0
self._act_info.bwd_peak_prev_cache_mem = 0
def _comp_leaf_model_info_impl(self):
- self._model_info.weight_bytes = 0
- self._model_info.grad_bytes = 0
- self._model_info.state_bytes = 0
+ self._model_info.dense_weight_bytes = 0
+ self._model_info.dense_grad_bytes = 0
+ self._model_info.dense_state_bytes = 0
def _comp_leaf_flops_info(self):
"""
@@ -807,7 +854,8 @@ def _comp_leaf_mem_accessed_info(self):
permutate1_mem_accessed = ( # none-fused: contiguous memory(drop_and_pad) or sort_chunks_by_idxs
2 * self.act_size_before_combined
) * self.dtype_to_element_size[self.strategy.dtype]
- permutate2_and_combine_mem_accessed = ( # fused: combine permuted_features by probs and scatter_add
+
+ permutate2_and_combine_mem_accessed = ( # fused-op: combine permuted_features by probs and scatter_add
self.act_size_before_combined + self.act_size_after_combined
) * self.dtype_to_element_size[self.strategy.dtype]
@@ -847,7 +895,8 @@ def __init__(
enable_recompute: bool,
mode:str,
strategy: StrategyConfig,
- system: SystemConfig
+ system: SystemConfig,
+ is_last_recompute: bool = False
) -> None:
super().__init__(local_expert_num, input_size, output_size, strategy, system)
assert mode in ['parallel', 'serial']
@@ -859,7 +908,10 @@ def __init__(
self.use_bias = use_bias # for now unless
self.has_cached_inputs = has_cached_inputs
self.enable_recompute = enable_recompute
-
+ self.is_last_recompute = is_last_recompute
+
+ if self.is_last_recompute and self.enable_recompute:
+ self.set_variance_node(True)
if self.strategy.fp8:
self.w_dtype = "fp8"
self.a_dtype = "fp8"
@@ -929,11 +981,10 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- token_num = self.output_info.tensors[0].size(0)
+ token_num = self.output_info_.tensors[0].size(0)
return token_num * self.output_size
- @property
- def output_info(self):
+ def create_output_info(self):
token_num = self.input_info.tensors[0].size(0)
origin_input_info = self.input_info.tensors[1:]
output_info = InputOutputInfo(
@@ -955,7 +1006,7 @@ def _comp_leaf_act_info_impl(self):
self._act_info.activation_mem_cache = (
self.micro_hidden_state_size * self.a_element_size # fp8
)
- if self.has_cached_inputs:
+ if self.has_cached_inputs or self.offload_inputs:
self._act_info.activation_mem_cache = 0
weight_size = (
self.local_expert_num
@@ -973,11 +1024,12 @@ def _comp_leaf_act_info_impl(self):
output_size = self.micro_output_grad_size * self.element_size # bf16
self._act_info.fwd_peak_mem_no_cache = input_size + output_size + (0 if self.strategy.use_accm_weight else weight_size)
self._act_info.fwd_peak_prev_cache_mem = 0
- self._act_info.bwd_peak_mem_no_cache = input_size + output_size + (grad_size if self.strategy.fp8 else 0)
+ self._act_info.bwd_peak_mem_no_cache = input_size + output_size + (grad_size if self.strategy.fp8 else 0) + (input_size if self.offload_inputs else 0)
self._act_info.bwd_peak_prev_cache_mem = 0
def _comp_leaf_model_info_impl(self):
weight_numel = self.local_expert_num * self.input_size * self.output_size
+ self._model_info.moe_weight_numel = weight_numel * self.strategy.ep_size * self.strategy.etp_size # Statistics the parameters of all etp ranks and ep ranks
self._model_info.moe_weight_bytes = weight_numel * self.w_element_size # fp8
self._model_info.moe_grad_bytes = (
self._model_info.moe_weight_bytes
@@ -1069,6 +1121,7 @@ def __init__(
mode:str,
strategy: StrategyConfig,
system: SystemConfig,
+ is_last_recompute: bool = False
) -> None:
super().__init__(local_expert_num, input_size, output_size, strategy, system)
assert mode in ['parallel', 'serial']
@@ -1080,6 +1133,9 @@ def __init__(
self.use_bias = use_bias
self.has_cached_inputs = has_cached_inputs
self.enable_recompute = enable_recompute
+ self.is_last_recompute = is_last_recompute
+ if self.is_last_recompute and self.enable_recompute:
+ self.set_variance_node(True)
if self.strategy.fp8:
self.w_dtype = "fp8"
self.a_dtype = "fp8"
@@ -1155,13 +1211,12 @@ def micro_hidden_state_size(self):
@property
def micro_output_grad_size(self):
# [B, S, H]
- token_num = self.output_info.tensors[0].size(0)
- hidden_size = self.output_info.tensors[0].size(1)
+ token_num = self.output_info_.tensors[0].size(0)
+ hidden_size = self.output_info_.tensors[0].size(1)
# hidden_size = self.output_info.tensors[0].size(2)
return token_num * hidden_size
- @property
- def output_info(self):
+ def create_output_info(self):
token_num = self.input_info.tensors[0].size(0)
origin_input_info = self.input_info.tensors[1:]
@@ -1206,6 +1261,7 @@ def _comp_leaf_act_info_impl(self):
def _comp_leaf_model_info_impl(self):
weight_numel = self.input_size * self.output_size * self.local_expert_num
+ self._model_info.moe_weight_numel = weight_numel * self.strategy.ep_size * self.strategy.etp_size # Statistics the parameters of all etp ranks and ep ranks
self._model_info.moe_weight_bytes = weight_numel * self.w_element_size # fp8
self._model_info.moe_grad_bytes = (
self._model_info.moe_weight_bytes
@@ -1293,11 +1349,16 @@ def __init__(self,
enable_recompute: bool,
mode:str,
strategy: StrategyConfig,
- system: SystemConfig
+ system: SystemConfig,
+ is_last_recompute: bool = False
):
super().__init__(strategy, system)
-
- self.quantizer = Float8Quantizer(enable_recompute=enable_recompute, strategy=strategy, system=system)
+ quantizer_recompute = False if strategy.cache_groupgemm_col_fp8_inputs else enable_recompute
+ self.quantizer = Float8Quantizer(enable_recompute=quantizer_recompute, strategy=strategy, system=system)
+ enable_cahce_bf16_inputs = not self.strategy.cache_groupgemm_col_fp8_inputs
+ if enable_cahce_bf16_inputs:
+ self.quantizer.offload_inputs = self.strategy.offload_groupgemm_col_inputs # the quantizer can perform offload When the input of bf16 needs to be cached
+
self.linear = GroupLinearCol(
layer_idx,
input_size,
@@ -1308,7 +1369,8 @@ def __init__(self,
enable_recompute,
mode,
strategy,
- system
+ system,
+ is_last_recompute
)
def forward(self, hidden_states, path_debug_context=None):
hidden_states = self.quantizer(hidden_states, path_debug_context)
@@ -1327,9 +1389,11 @@ def __init__(self,
enable_recompute: bool,
mode:str,
strategy: StrategyConfig,
- system: SystemConfig,):
+ system: SystemConfig,
+ if_first_recompute: bool = False,
+ is_last_recompute: bool = False
+ ):
super().__init__(strategy, system)
-
self.quantizer = Float8Quantizer(enable_recompute=enable_recompute, strategy=strategy, system=system)
self.linear = GroupLinearRow(
layer_idx,
@@ -1342,6 +1406,7 @@ def __init__(self,
mode,
strategy,
system,
+ is_last_recompute
)
def forward(self, hidden_states, path_debug_context=None):
@@ -1376,14 +1441,14 @@ def __init__(self,
if self.config.use_swiglu
else ffn_hidden_size
)
-
+ self.mlp_recompute = mlp_recompute
self.shared_expert = None
if getattr(self.config, "moe_shared_expert_intermediate_size", None) is not None:
self.shared_expert = MLP(
layer_idx=f"{layer_idx}-shareExpert",
config=self.config,
enable_recompute=enable_recompute, # for old version
- mlp_recompute=mlp_recompute,
+ mlp_recompute_conf=mlp_recompute,
strategy=strategy,
system=system,
intermediate_size=self.config.moe_shared_expert_intermediate_size
@@ -1391,6 +1456,7 @@ def __init__(self,
GroupLinearCol_ = QuantizedGroupLinearCol if self.strategy.fp8 else GroupLinearCol
GroupLinearRow_ = QuantizedGroupLinearRow if self.strategy.fp8 else GroupLinearRow
+
self.router = Router(
layer_idx=layer_idx,
hidden_size=self.config.hidden_size,
@@ -1427,6 +1493,15 @@ def __init__(self,
strategy=strategy,
system=system,
)
+ if self.strategy.fp8:
+ # fp8
+ if self.strategy.cache_groupgemm_col_fp8_inputs:
+ self.group_linear1.linear.offload_inputs = self.strategy.offload_groupgemm_col_inputs
+ else:
+ self.group_linear1.quantizer.offload_inputs = self.strategy.offload_groupgemm_col_inputs
+ else:
+ # bf16
+ self.group_linear1.offload_inputs = self.strategy.offload_groupgemm_col_inputs
if self.config.use_swiglu:
self.expert_activation_layer = Swiglu(
@@ -1435,6 +1510,7 @@ def __init__(self,
enable_recompute=mlp_recompute.linear_recompute,
strategy=strategy,
system=system,
+ is_weighted_silu= self.strategy.dispatch_probs
)
else:
self.expert_activation_layer =Gelu(
@@ -1450,6 +1526,7 @@ def __init__(self,
local_expert_num=self.local_expert_num,
has_cached_inputs=False,
enable_recompute=mlp_recompute.linear_recompute,
+ is_last_recompute = True,
mode=self.config.group_linear_mode,
use_bias=False,
strategy=strategy,
@@ -1474,30 +1551,42 @@ def __init__(self,
mlp_recompute.linear_recompute and
self.shared_expert.recompute_granularity == "full" if self.shared_expert else True):
self.recompute_granularity = "submodule"
-
+
+ def preprocess(self, input_info:InputOutputInfo):
+ self.unpermutation.set_ori_shape(input_info.tensors[0].shape.copy())
+
def forward(self, input_info:InputOutputInfo, path_debug_context:PathDebugContext):
+ self.preprocess(input_info)
if self.shared_expert:
shared_out = self.shared_expert(input_info, path_debug_context)
- hidden_states = self.router(input_info, path_debug_context) # add router scores
-
- hidden_states = self.permutation(hidden_states, path_debug_context)
+ probs = self.router(input_info, path_debug_context) # add router scores
- hidden_states = self.group_linear1(hidden_states, path_debug_context)
- hidden_states = self.expert_activation_layer(hidden_states, path_debug_context)
- hidden_states = self.group_linear2(hidden_states, path_debug_context)
-
- merge_input = InputOutputInfo(
- tensors=[
- hidden_states,
- *input_info.tensors,
- ]
- )
- out = self.unpermutation(merge_input, path_debug_context)
- # print(f'--------- out={out}')
+ if self.strategy.dispatch_probs:
+ permute_hidden_states = self.permutation(Input(tensors=[input_info.tensors[0],
+ probs.tensors[0]]),
+ path_debug_context)
+ grou1_out = self.group_linear1(permute_hidden_states, path_debug_context)
+ act_out = self.expert_activation_layer(Input(tensors=[grou1_out.tensors[0], probs.tensors[0]]), path_debug_context)
+ group2_out = self.group_linear2(act_out, path_debug_context)
+ out = self.unpermutation(group2_out, path_debug_context)
+ else:
+ permute_hidden_states = self.permutation(Input(tensors=[input_info.tensors[0],
+ probs.tensors[0]]),path_debug_context) # pass probs to Permutation to cache.
+ grou1_out = self.group_linear1(permute_hidden_states, path_debug_context)
+ act_out = self.expert_activation_layer(grou1_out, path_debug_context)
+ group2_out = self.group_linear2(act_out, path_debug_context)
+ out = self.unpermutation(Input(tensors=[group2_out.tensors[0], probs.tensors[0]]), path_debug_context)
+
# FIXME(sherry): add mul, routed_expert hidden_states * router scores
if self.shared_expert:
- return out + shared_out
+ # return out + shared_out
+ return AddFunction.apply(parent_model=self,
+ enable_recompute=self.recompute_granularity == 'full_block',
+ tensor_size1=out,
+ tensor_size2=shared_out,
+ path_debug_context=path_debug_context,
+ name='SharedExpertAddFunction')
return out
def prefill(self, args, call_stk='', com_buff=None):
diff --git a/simumax/core/transformer/simu_ops.py b/simumax/core/transformer/simu_ops.py
index 5042017..39fda28 100644
--- a/simumax/core/transformer/simu_ops.py
+++ b/simumax/core/transformer/simu_ops.py
@@ -9,7 +9,7 @@ def split(tensor_size: TensorSize, split_size_or_sections:List[int], dim = -1):
split_size_or_sections = [tensor_size[dim] // split_size_or_sections] * split_size_or_sections
assert tensor_size[dim] == sum(split_size_or_sections), f"tensor_size[dim]={tensor_size[dim]} sum(split_size_or_sections)={sum(split_size_or_sections)}"
- return [tensor_size.new(dim, size) for size in split_size_or_sections]
+ return [tensor_size.new_with_dim(dim, size) for size in split_size_or_sections]
def transpose(tensor_size: TensorSize, dim0: int, dim1: int):
@@ -20,7 +20,7 @@ def cat(tensor_sizes: List[TensorSize], dim:int = -1):
if len(tensor_sizes) == 0:
return tensor_sizes
concat_size = sum([t[dim] for t in tensor_sizes])
- return tensor_sizes[0].new(dim, concat_size)
+ return tensor_sizes[0].new_with_dim(dim, concat_size)
def unsqueeze(tensor_size: TensorSize, dim:int):
return tensor_size.unsqeeze(dim)
diff --git a/simumax/core/utils.py b/simumax/core/utils.py
index 8584eed..bb44f58 100644
--- a/simumax/core/utils.py
+++ b/simumax/core/utils.py
@@ -3,6 +3,20 @@
import json
import os, subprocess
+def wrap_name(src):
+ return f"_orig_{src}"
+
+def add_attr(module, name, target):
+ setattr(module, name, target)
+
+def wrap_attr(module, name, wrapper):
+ target = getattr(module, name)
+ setattr(module, wrap_name(name), target)
+ setattr(module, name, wrapper)
+
+def replace_attr(module, name, target):
+ wrap_attr(module, name, target)
+
class HumanReadableSize:
"""Convert a size in bytes to a human-readable format."""
diff --git a/simumax/tuning/strategy_searcher.py b/simumax/tuning/strategy_searcher.py
index dd2f830..d6872e7 100644
--- a/simumax/tuning/strategy_searcher.py
+++ b/simumax/tuning/strategy_searcher.py
@@ -20,14 +20,12 @@
"interleaving_size": [1],
"zero_state": [1],
"use_fused_norm": [True],
- "no_sync": [True],
"use_math_sdp": [False],
"use_flash_sdp": [True],
"use_fp32_accum_grad": [True],
"use_fused_swiglu": [True],
"enable_recompute": [False, True],
"recompute_granularity": ["full_block"],
- "skip_ckpt_micro_batch_num": [0],
"mem_factor": [0.94],
}
diff --git a/simumax/utils.py b/simumax/utils.py
index 86f62cd..d9df08e 100644
--- a/simumax/utils.py
+++ b/simumax/utils.py
@@ -31,7 +31,7 @@ def get_config(key, version, r_maps:dict, d_maps:dict, m_type):
assert key in r_maps.keys(), f"{key} not found, please add {m_type} config in {r_maps['root']}"
return r_maps[key]
elif version == 'dev':
- assert key in DEV_MODELS.keys(), f"{key} not found, please add {m_type} config in {d_maps['root']}"
+ assert key in d_maps.keys(), f"{key} not found, please add {m_type} config in {d_maps['root']}"
return d_maps[key]
else:
raise ValueError('type must be release or dev')
@@ -76,4 +76,44 @@ def show_simu_system(version='release'):
raise ValueError('type must be release or dev')
+def parse_parallel_config(s):
+ parts = s.split('_')
+ config = {}
+ for part in parts:
+ if part.startswith('ep'):
+ config['ep_size'] = int(part[2:])
+ elif part.startswith('tp'):
+ config['tp_size'] = int(part[2:])
+ elif part.startswith('pp'):
+ config['pp_size'] = int(part[2:])
+ return config
+
+def create_default_strategy(strs="", seq_len=4096, micro_batch_size=1, micro_batch_num=81, dtype='bf16'):
+ default_config = {
+ "seq_len": seq_len,
+ "micro_batch_size": micro_batch_size,
+ "micro_batch_num": micro_batch_num,
+ "dtype": "bf16",
+ "world_size": 8,
+ "tp_size": 1,
+ "pp_size": 1,
+ "ep_size": 1,
+ "etp_size": 1,
+ "moe_dispatcher_policy": "all2all",
+ "enable_sequence_parallel": True,
+ "interleaving_size": 1,
+ "zero_state": 1,
+ "enable_dropout": False,
+ "use_fused_norm": True,
+ "use_math_sdp": False,
+ "use_flash_sdp": True,
+ "use_fp32_accum_grad": True,
+ "enable_recompute": True,
+ "mem_factor": 0.94,
+ "fp8" : True if dtype == 'fp8' else False,
+ }
+ default_config.update(parse_parallel_config(strs))
+ return default_config
+
+
\ No newline at end of file