mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-03 08:14:29 +08:00
h-guo18
82d96a635f
## What does this PR do?
Refactors EAGLE3 training to use a single base YAML config with
OmegaConf dotlist overrides.
**Type of change:** Refactor
## Changes
- Single base config
`modelopt_recipes/speculative_decoding/_base_eagle3.yaml` for all EAGLE3
training; removed per-model child YAMLs.
- `launch_train.sh` accepts `--config <yaml>` plus dotlist overrides
(e.g. `model.model_name_or_path=xxx`).
- Removed `__base__` YAML inheritance logic from `main.py`.
- `dp_shard_size` default changed from `0` sentinel to `None` for
clarity.
- Removed `eagle_config.json` and `fsdp_config.json`; architecture
config is now nested under `eagle.eagle_architecture_config` in YAML.
- `train_eagle3_and_export.sh` now uses base YAML + dotlist instead of
generating a temporary YAML.
- Updated README and tests accordingly.
## Usage
```bash
# Online training
./launch_train.sh \
--config ../../modelopt_recipes/speculative_decoding/_base_eagle3.yaml \
model.model_name_or_path=meta-llama/Llama-3.2-1B \
data.data_path=input_conversations/train.jsonl \
training.output_dir=ckpts/llama-3.2-1b-online
# Offline training
./launch_train.sh \
--config ../../modelopt_recipes/speculative_decoding/_base_eagle3.yaml \
model.model_name_or_path=meta-llama/Llama-3.2-1B \
data.offline_data_path=$HIDDEN_STATES_DIR \
training.output_dir=ckpts/llama-3.2-1b-offline
```
---------
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
76 lines
3.0 KiB
Bash
Executable File
76 lines
3.0 KiB
Bash
Executable File
#!/bin/bash
|
|
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# Usage:
|
|
# Single GPU: ./launch_train.sh --config ../../modelopt_recipes/general/speculative_decoding/eagle3.yaml model.model_name_or_path=xxx
|
|
# Multi-node: ./launch_train.sh --config ../../modelopt_recipes/general/speculative_decoding/eagle3.yaml --num_nodes 2 --head_node_ip <IP>
|
|
# With overrides: ./launch_train.sh --config my.yaml model.model_name_or_path=xxx training.output_dir=yyy
|
|
#
|
|
# Extra key=value args are forwarded as OmegaConf dotlist overrides to main.py.
|
|
# All training config (model, data, hyperparams, eagle, fsdp) lives in the YAML file.
|
|
# Only multi-node routing args are passed here; mixed_precision is fixed to bf16.
|
|
|
|
set -eo pipefail
|
|
|
|
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
|
|
|
|
CONFIG_FILE=""
|
|
NUM_NODES=1
|
|
HEAD_NODE_IP=""
|
|
EXTRA_ARGS=()
|
|
while [ $# -gt 0 ]; do
|
|
case "$1" in
|
|
--config*) if [[ "$1" != *=* ]]; then shift; fi; CONFIG_FILE="${1#*=}" ;;
|
|
--num_nodes*) if [[ "$1" != *=* ]]; then shift; fi; NUM_NODES="${1#*=}" ;;
|
|
--head_node_ip*) if [[ "$1" != *=* ]]; then shift; fi; HEAD_NODE_IP="${1#*=}" ;;
|
|
*) EXTRA_ARGS+=("$1") ;;
|
|
esac
|
|
shift
|
|
done
|
|
|
|
if [ -z "$CONFIG_FILE" ]; then
|
|
>&2 echo "Usage: ./launch_train.sh --config <yaml_file> [--num_nodes N] [--head_node_ip IP] [key=value ...]"
|
|
exit 1
|
|
fi
|
|
|
|
# GPU count detection
|
|
if [[ "$NUM_NODES" != "1" ]]; then
|
|
GPU_PER_NODE=${GPU_PER_NODE:-$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)}
|
|
TOTAL_GPU=$((NUM_NODES * GPU_PER_NODE))
|
|
echo "Total GPUs: $TOTAL_GPU (NUM_NODES: $NUM_NODES, GPU_PER_NODE: $GPU_PER_NODE)"
|
|
else
|
|
TOTAL_GPU=$(python3 -c "import torch; print(torch.cuda.device_count())")
|
|
echo "Total GPUs: $TOTAL_GPU (single node)"
|
|
fi
|
|
|
|
# Multi-node routing args (accelerate only; training config comes from the YAML)
|
|
MULTI_NODE_ARGS=""
|
|
if [[ "$NUM_NODES" != "1" ]]; then
|
|
MULTI_NODE_ARGS="--num_processes $TOTAL_GPU \
|
|
--num_machines $NUM_NODES \
|
|
--machine_rank $SLURM_PROCID \
|
|
--rdzv_backend c10d \
|
|
--main_process_ip $HEAD_NODE_IP \
|
|
--main_process_port 29500"
|
|
fi
|
|
|
|
export TOKENIZERS_PARALLELISM=False
|
|
|
|
set -x
|
|
start_time=$(date +%s)
|
|
sh -c "accelerate launch --mixed_precision bf16 $MULTI_NODE_ARGS ${SCRIPT_DIR}/main.py --config $CONFIG_FILE ${EXTRA_ARGS[*]}"
|
|
echo "Total time: $(( $(date +%s) - $start_time )) seconds"
|