【训练营】小模型训练支持

CMakeLists.txt

@@ -104,7 +104,7 @@ if(USE_CUDA)
   file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu)
 
   add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS})
-  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90")
+  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;86;89;90")
 
   target_link_libraries(infini_train_cuda_kernels
     PUBLIC

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,210 @@ 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
+
+Train an MLP classifier (the original example) or a CNN classifier on MNIST.
+The CNN stacks `Conv2d(1, 16, 3) -> ReLU -> Conv2d(16, 32, 3) -> ReLU ->
+Flatten -> Linear(18432, 10)` and reports test loss and accuracy after every
+epoch.
+
+```bash
+# MLP (default: --model mlp)
+./build/mnist \
+  --device cpu \
+  --dataset data/mnist
+
+# CNN on CUDA
+./build/mnist \
+  --model cnn \
+  --device cuda \
+  --dataset data/mnist
 ```
 
-### 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
-```
+Pass `--init_weights <checkpoint_dir>` to load initial weights from an
+InfiniTrain checkpoint (`model.ckpt` inside the directory), for example one
+exported by PyTorch for numerical-alignment runs.
+
+```bash
+./build/mnist \
+  --model cnn \
+  --device cuda \
+  --dataset data/mnist \
+  --init_weights data/cnn_align
+```
+
+Pass `--metrics_file <path>` to append per-step training losses and per-epoch
+test loss / accuracy as JSON lines, which can be uploaded to
+[SwanLab](https://swanlab.cn) for training visualization:
+
+```bash
+./build/mnist --model cnn --device cuda --dataset data/mnist --metrics_file metrics.jsonl
+SWANLAB_API_KEY=<key> python3 scripts/swanlab_upload.py \
+  --metrics metrics.jsonl --name cnn-cuda-3epoch-lr0.05 --model cnn --device cuda --lr 0.05
+```
+
+##### 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
+```
+
+The MNIST example supports the same DDP launch (`--device cuda` is required;
+each rank trains on its own GPU and logs only on rank 0):
+
+```bash
+./build/infini_run \
+  --nnodes=1 \
+  --nproc_per_node=2 \
+  ./build/mnist \
+    --model cnn \
+    --device cuda \
+    --dataset data/mnist \
+    --num_epoch 3 \
+    --lr 0.1 \
+    --metrics_file metrics_ddp.jsonl
+```
+
+#### 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)
 
@@ -316,4 +366,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.