【训练营】Pipeline 并行自定义布局

README.md

@@ -50,7 +50,7 @@ Build Options:
 
 > Both options are optional and can be disabled for CPU-only builds.
 
-## ✨ InfiniTrain Overview
+## ✨ InfiniTrain Overview
 
 ### ✔ Support Matrix
 
@@ -96,160 +96,160 @@ For example, the `llama3` example produces a binary named `llama3`.
 
 To view available runtime options:
 
-```bash
-./build/llama3 --help
+```bash
+./build/llama3 --help
+```
+
+### Getting Started
+
+#### Prepare Datasets and Weights
+
+Run the asset preparation script from the repository root. Prepared files are
+written to `data/` by default.
+
+```bash
+# MNIST dataset
+./scripts/assets/prepare-infinitrain-assets.sh mnist
+
+# GPT-2 124M weights, tokenizer, and tokenized TinyShakespeare data
+./scripts/assets/prepare-infinitrain-assets.sh gpt2
+
+# LLaMA 3.2 1B weights and tokenized TinyShakespeare data
+HF_TOKEN=hf_xxx ./scripts/assets/prepare-infinitrain-assets.sh llama3
+```
+
+Preparing LLaMA requires access to the gated
+`meta-llama/Llama-3.2-1B` repository. Accept its license on Hugging Face and
+provide `HF_TOKEN`, or authenticate with `hf auth login`, before running the
+command. The complete LLaMA preparation requires approximately 8.5 GB of free
+disk space, including the downloaded checkpoint and converted FP32 weights.
+
+Use `DATA_DIR` to write the assets elsewhere, or prepare all supported assets
+in one invocation:
+
+```bash
+DATA_DIR=/path/to/data \
+HF_TOKEN=hf_xxx \
+./scripts/assets/prepare-infinitrain-assets.sh all
+```
+
+#### Model Examples
+
+The generated files can be passed directly to the corresponding executables:
+
+##### MNIST
+
+```bash
+./build/mnist \
+  --device cpu \
+  --dataset data/mnist
+```
+
+##### GPT-2 124M
+
+```bash
+./build/gpt2 \
+  --device cuda \
+  --input_bin data/gpt2/tiny_shakespeare_train.bin \
+  --input_val_bin data/gpt2/tiny_shakespeare_val.bin \
+  --tokenizer_bin data/gpt2/gpt2_tokenizer.bin \
+  --llmc_filepath data/gpt2/gpt2_124M.bin \
+  --num_iteration 10
+```
+
+##### LLaMA 3.2 1B
+
+```bash
+./build/llama3 \
+  --device cuda \
+  --input_bin data/llama3/tiny_shakespeare_train.bin \
+  --input_val_bin data/llama3/tiny_shakespeare_val.bin \
+  --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
+  --num_iteration 10
 ```
 
-### Getting Started
-
-#### Prepare Datasets and Weights
-
-Run the asset preparation script from the repository root. Prepared files are
-written to `data/` by default.
-
-```bash
-# MNIST dataset
-./scripts/assets/prepare-infinitrain-assets.sh mnist
-
-# GPT-2 124M weights, tokenizer, and tokenized TinyShakespeare data
-./scripts/assets/prepare-infinitrain-assets.sh gpt2
-
-# LLaMA 3.2 1B weights and tokenized TinyShakespeare data
-HF_TOKEN=hf_xxx ./scripts/assets/prepare-infinitrain-assets.sh llama3
-```
-
-Preparing LLaMA requires access to the gated
-`meta-llama/Llama-3.2-1B` repository. Accept its license on Hugging Face and
-provide `HF_TOKEN`, or authenticate with `hf auth login`, before running the
-command. The complete LLaMA preparation requires approximately 8.5 GB of free
-disk space, including the downloaded checkpoint and converted FP32 weights.
-
-Use `DATA_DIR` to write the assets elsewhere, or prepare all supported assets
-in one invocation:
-
-```bash
-DATA_DIR=/path/to/data \
-HF_TOKEN=hf_xxx \
-./scripts/assets/prepare-infinitrain-assets.sh all
-```
-
-#### Model Examples
-
-The generated files can be passed directly to the corresponding executables:
-
-##### MNIST
-
-```bash
-./build/mnist \
-  --device cpu \
-  --dataset data/mnist
-```
-
-##### GPT-2 124M
-
-```bash
-./build/gpt2 \
-  --device cuda \
-  --input_bin data/gpt2/tiny_shakespeare_train.bin \
-  --input_val_bin data/gpt2/tiny_shakespeare_val.bin \
-  --tokenizer_bin data/gpt2/gpt2_tokenizer.bin \
-  --llmc_filepath data/gpt2/gpt2_124M.bin \
-  --num_iteration 10
-```
-
-##### LLaMA 3.2 1B
-
-```bash
-./build/llama3 \
-  --device cuda \
-  --input_bin data/llama3/tiny_shakespeare_train.bin \
-  --input_val_bin data/llama3/tiny_shakespeare_val.bin \
-  --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
-  --num_iteration 10
-```
-
-### Launch Modes
-
-GPT-2 and LLaMA training support both thread-based and process-based launches.
-The examples below use LLaMA, but the same launch modes also apply to GPT-2.
-
-#### Direct Launch
-
-Running a model executable directly uses one process and one device by default.
-Set `--nthread_per_process` to use multiple execution threads and devices in the
-same process:
-
-```bash
-./build/llama3 \
-  --device cuda \
-  --input_bin data/llama3/tiny_shakespeare_train.bin \
-  --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
-  --nthread_per_process 8 \
-  --num_iteration 10
-```
-
-#### Single-node Multi-process Launch
-
-Use `infini_run` to start multiple training processes on one node. Each process
-uses one execution thread by default:
-
-```bash
-./build/infini_run \
-  --nnodes=1 \
-  --nproc_per_node=8 \
-  ./build/llama3 \
-    --device cuda \
-    --input_bin data/llama3/tiny_shakespeare_train.bin \
-    --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
-    --num_iteration 10
-```
-
-#### Multi-node Multi-process Launch
-
-Run the following command on every node with the same rendezvous settings and
-a distinct `node_rank`:
-
-```bash
-./build/infini_run \
-  --nnodes=2 \
-  --nproc_per_node=4 \
-  --node_rank=[rank_id] \
-  --rdzv_endpoint=[master_addr]:29500 \
-  --rdzv_id=[job_id] \
-  ./build/llama3 \
-    --device cuda \
-    --input_bin data/llama3/tiny_shakespeare_train.bin \
-    --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
-    --num_iteration 10 \
-    --tensor_parallel 2 \
-    --pipeline_parallel 2 \
-    --sequence_parallel
-```
-
-`--nproc_per_node` and `--nthread_per_process` can be combined. The total
-training world size is:
-
-```text
-world_size = nnodes × nproc_per_node × nthread_per_process
-```
+### Launch Modes
+
+GPT-2 and LLaMA training support both thread-based and process-based launches.
+The examples below use LLaMA, but the same launch modes also apply to GPT-2.
+
+#### Direct Launch
+
+Running a model executable directly uses one process and one device by default.
+Set `--nthread_per_process` to use multiple execution threads and devices in the
+same process:
+
+```bash
+./build/llama3 \
+  --device cuda \
+  --input_bin data/llama3/tiny_shakespeare_train.bin \
+  --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
+  --nthread_per_process 8 \
+  --num_iteration 10
+```
+
+#### Single-node Multi-process Launch
+
+Use `infini_run` to start multiple training processes on one node. Each process
+uses one execution thread by default:
+
+```bash
+./build/infini_run \
+  --nnodes=1 \
+  --nproc_per_node=8 \
+  ./build/llama3 \
+    --device cuda \
+    --input_bin data/llama3/tiny_shakespeare_train.bin \
+    --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
+    --num_iteration 10
+```
+
+#### Multi-node Multi-process Launch
+
+Run the following command on every node with the same rendezvous settings and
+a distinct `node_rank`:
+
+```bash
+./build/infini_run \
+  --nnodes=2 \
+  --nproc_per_node=4 \
+  --node_rank=[rank_id] \
+  --rdzv_endpoint=[master_addr]:29500 \
+  --rdzv_id=[job_id] \
+  ./build/llama3 \
+    --device cuda \
+    --input_bin data/llama3/tiny_shakespeare_train.bin \
+    --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \
+    --num_iteration 10 \
+    --tensor_parallel 2 \
+    --pipeline_parallel 2 \
+    --sequence_parallel
+```
+
+`--nproc_per_node` and `--nthread_per_process` can be combined. The total
+training world size is:
+
+```text
+world_size = nnodes × nproc_per_node × nthread_per_process
+```
 
 ### Parallelism Strategies
 
-#### Distributed Data Parallelism (DDP)
-
-For a direct launch with TP and PP disabled, the following starts eight
-data-parallel workers in one process:
-
-```bash
---nthread_per_process 8  # 8-way DDP when TP=1 and PP=1
-```
-
-For all launch modes, the data-parallel size is derived from the total world
-size after accounting for tensor and pipeline parallelism:
-
-```text
-data_parallel_size = world_size / (tensor_parallel × pipeline_parallel)
-```
+#### Distributed Data Parallelism (DDP)
+
+For a direct launch with TP and PP disabled, the following starts eight
+data-parallel workers in one process:
+
+```bash
+--nthread_per_process 8  # 8-way DDP when TP=1 and PP=1
+```
+
+For all launch modes, the data-parallel size is derived from the total world
+size after accounting for tensor and pipeline parallelism:
+
+```text
+data_parallel_size = world_size / (tensor_parallel × pipeline_parallel)
+```
 
 #### Tensor Parallelism (TP)
 
@@ -263,8 +263,15 @@ data_parallel_size = world_size / (tensor_parallel × pipeline_parallel)
 ```bash
 --pipeline_parallel 8     		# 8 pipeline stages
 --virtual_pipeline_parallel 4  	# Virtual pipeline for better load balancing
+--pipeline_layer_partition 4,8,6,6  # Optional custom per-stage layer counts
 ```
 
+`--pipeline_layer_partition` lets you assign a non-uniform but contiguous number
+of transformer layers to each stage (e.g. `4,8,6,6` for 4 stages / 24 layers).
+The sum of entries must equal the model layer count. When omitted, the default
+uniform partition (compatible with vPP) is used. See
+[`docs/pipeline_layout_guide.md`](docs/pipeline_layout_guide.md) for details.
+
 #### Combining Parallelism Strategies
 
 Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scale training across devices and nodes.
@@ -316,4 +323,4 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal
    optimizations.
 
   Integrated a CTest + GTest based testing infrastructure to strengthen the
-   framework's automated test workflow.
+   framework's automated test workflow.