wehub-resource-sync
305 lines
12 KiB
Python
305 lines
12 KiB
Python
"""
|
|
Simplified LLM Fine-tuning Configuration Validator
|
|
|
|
Two-step validation:
|
|
1. Parameter filtering - Remove unsupported parameters
|
|
2. Micro-batch testing - Runtime validation with small dataset
|
|
"""
|
|
|
|
import json
|
|
import re
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Set
|
|
|
|
import yaml
|
|
|
|
from rdagent.components.coder.finetune.conf import (
|
|
FT_DEBUG_YAML_FILE_NAME,
|
|
FT_TEST_PARAMS_FILE_NAME,
|
|
get_ft_env,
|
|
get_workspace_prefix,
|
|
)
|
|
from rdagent.core.experiment import FBWorkspace
|
|
from rdagent.log import rdagent_logger as logger
|
|
from rdagent.scenarios.finetune.scen.llama_factory_manager import LLaMAFactory_manager
|
|
|
|
DIRNAME = Path(__file__).absolute().resolve().parent
|
|
|
|
# System-managed parameters that are automatically injected during validation.
|
|
# These should NOT be checked for alignment in eval prompts.
|
|
# Single source of truth: modify here to change injected parameters.
|
|
SYSTEM_MANAGED_PARAMS = {
|
|
"overwrite_cache": True, # Avoid HF datasets cache lock contention
|
|
"save_only_model": True, # Save disk space
|
|
# "save_total_limit": 1, # Limit checkpoint count to save disk space
|
|
"output_dir": "./output", # Standardize model output location
|
|
"per_device_eval_batch_size": 1, # Prevent OOM during evaluation
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ValidationResult:
|
|
"""Configuration validation result"""
|
|
|
|
success: bool
|
|
filtered_config: str
|
|
execution_output: str = "" # Parsed/summarized output for LLM
|
|
raw_stdout: str = "" # Full raw stdout for UI display
|
|
errors: List[str] = field(default_factory=list)
|
|
execution_time: float = 0.0
|
|
|
|
|
|
class LLMConfigValidator:
|
|
"""LLM configuration validator with two-step validation:
|
|
|
|
1. Parameter filtering - Remove unsupported parameters
|
|
2. Micro-batch test - Runtime validation with small dataset
|
|
|
|
The micro-batch test inherently validates completeness, so no separate completeness check is needed.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._supported_params_cache: Optional[Set[str]] = None
|
|
|
|
def validate_and_test(self, config_yaml: str, workspace: FBWorkspace, env) -> ValidationResult:
|
|
"""Three-step validation: parameter filtering + injection + micro-batch testing"""
|
|
start_time = time.time()
|
|
|
|
# Step 1: Parameter filtering
|
|
filtered_config, removed_params = self._filter_parameters(config_yaml)
|
|
|
|
# Step 2: Inject required parameters for multi-task environments
|
|
injected_config = self._inject_required_parameters(filtered_config)
|
|
|
|
# Step 3: Micro-batch testing (validates everything at runtime)
|
|
result = self._run_micro_batch_test(injected_config, workspace, env)
|
|
result.execution_time = time.time() - start_time
|
|
|
|
# Add filtered params info to execution_output for agent learning
|
|
if removed_params:
|
|
filter_info = (
|
|
f"\n\n[Filtered Parameters] {len(removed_params)} unsupported params removed: {removed_params}"
|
|
)
|
|
result.execution_output += filter_info
|
|
|
|
return result
|
|
|
|
def _filter_parameters(self, config_yaml: str) -> tuple[str, List[str]]:
|
|
"""Filter configuration parameters to only include supported ones.
|
|
|
|
Returns:
|
|
tuple: (filtered_yaml, removed_params_list)
|
|
"""
|
|
config_dict = yaml.safe_load(config_yaml)
|
|
if not isinstance(config_dict, dict):
|
|
return config_yaml, []
|
|
|
|
supported_params = self._get_supported_parameters()
|
|
|
|
filtered_config = {}
|
|
removed_params = []
|
|
for k, v in config_dict.items():
|
|
if k in supported_params:
|
|
filtered_config[k] = v
|
|
else:
|
|
removed_params.append(k)
|
|
|
|
if removed_params:
|
|
logger.info(f"Filtered out {len(removed_params)} unsupported parameters: {removed_params}")
|
|
|
|
return yaml.dump(filtered_config, default_flow_style=False, sort_keys=False), removed_params
|
|
|
|
def _inject_required_parameters(self, config_yaml: str) -> str:
|
|
"""Inject required parameters for multi-task environments.
|
|
|
|
Uses SYSTEM_MANAGED_PARAMS as the single source of truth.
|
|
"""
|
|
config = yaml.safe_load(config_yaml)
|
|
if not isinstance(config, dict):
|
|
return config_yaml
|
|
|
|
config.update(SYSTEM_MANAGED_PARAMS)
|
|
|
|
logger.info(f"Injected required parameters: {SYSTEM_MANAGED_PARAMS}")
|
|
return yaml.dump(config, default_flow_style=False, sort_keys=False)
|
|
|
|
def _get_supported_parameters(self) -> Set[str]:
|
|
"""Get supported parameters from LlamaFactory Manager"""
|
|
if self._supported_params_cache is not None:
|
|
return self._supported_params_cache
|
|
|
|
all_params = LLaMAFactory_manager.get_parameters()
|
|
|
|
# Extract all parameter names from all parameter types (including nested structures)
|
|
supported_params = set()
|
|
for param_type, params_dict in all_params.items():
|
|
if isinstance(params_dict, dict):
|
|
# Recursively extract parameter names from nested dictionaries
|
|
for key, value in params_dict.items():
|
|
if isinstance(value, dict) and "name" in value:
|
|
# This is a parameter definition with metadata
|
|
supported_params.add(key)
|
|
elif isinstance(value, dict):
|
|
# This is a nested category (e.g., BaseModelArguments, LoraArguments)
|
|
# Extract parameter names from the nested structure
|
|
for nested_key, nested_value in value.items():
|
|
if isinstance(nested_value, dict) and "name" in nested_value:
|
|
supported_params.add(nested_key)
|
|
|
|
if not supported_params:
|
|
raise RuntimeError("No parameters found in LlamaFactory Manager")
|
|
|
|
logger.info(f"Loaded {len(supported_params)} parameters from LlamaFactory Manager")
|
|
self._supported_params_cache = supported_params
|
|
return supported_params
|
|
|
|
def _parse_execution_log(self, stdout: str, exit_code: int, failed_stage: str = None) -> str:
|
|
"""Parse execution log and extract key information for LLM evaluation.
|
|
|
|
Reduces log from ~36k tokens to ~500 tokens by extracting only:
|
|
- Status and exit code
|
|
- Error messages (if any)
|
|
- Training metrics (if successful)
|
|
- Warnings (limited)
|
|
- Timeout and stage information (if applicable)
|
|
|
|
Args:
|
|
stdout: The execution output
|
|
exit_code: The process exit code
|
|
failed_stage: Which stage failed - "data_processing" or "training"
|
|
"""
|
|
result = {
|
|
"status": "success" if exit_code == 0 else "failed",
|
|
"exit_code": exit_code,
|
|
}
|
|
|
|
# Handle timeout (exit_code 124)
|
|
if exit_code == 124:
|
|
result["timeout"] = True
|
|
if failed_stage:
|
|
result["failed_stage"] = failed_stage
|
|
|
|
# 1. Extract error information (highest priority)
|
|
# Strategy: extract rank0's error block (each line prefixed with [rank0]:)
|
|
error_text = None
|
|
|
|
# Method A: Extract [rank0]: prefixed lines and reconstruct traceback
|
|
rank0_lines = re.findall(r"\[rank0\]:[^\n]+", stdout)
|
|
if rank0_lines:
|
|
rank0_block = "\n".join(line.replace("[rank0]: ", "").replace("[rank0]:", "") for line in rank0_lines)
|
|
# Find traceback in rank0 block
|
|
tb_match = re.search(
|
|
r"Traceback \(most recent call last\):.*?(?:Error|Exception):[^\n]+", rank0_block, re.DOTALL
|
|
)
|
|
if tb_match:
|
|
error_text = tb_match.group(0)
|
|
|
|
# Method B: Fallback to generic traceback (no rank prefix)
|
|
# Use findall to get ALL tracebacks, then keep the first one (root cause)
|
|
if not error_text:
|
|
all_tracebacks = re.findall(
|
|
r"Traceback \(most recent call last\):.*?(?:Error|Exception):[^\n]+", stdout, re.DOTALL
|
|
)
|
|
if all_tracebacks:
|
|
# First traceback is usually the root cause
|
|
error_text = all_tracebacks[0]
|
|
if len(all_tracebacks) > 1:
|
|
error_text += f"\n\n[Note: {len(all_tracebacks)} total errors, showing root cause]"
|
|
|
|
if error_text:
|
|
# Limit length but keep from the END (actual error type/message is at the end of traceback)
|
|
result["error"] = error_text[-4000:] if len(error_text) > 4000 else error_text
|
|
|
|
# 2. Extract training information
|
|
if "Running training" in stdout:
|
|
result["training_started"] = True
|
|
|
|
# Extract training config
|
|
# NOTE: we may have log like "Num examples = 1,000,000" and "Num Epochs = 1,000"; So we need to handle ","
|
|
num_examples = re.search(r"Num examples\s*=\s*([\d,]+)", stdout)
|
|
num_epochs = re.search(r"Num Epochs\s*=\s*([\d,]+)", stdout)
|
|
if num_examples:
|
|
result["num_examples"] = int(num_examples.group(1).replace(",", ""))
|
|
if num_epochs:
|
|
result["num_epochs"] = int(num_epochs.group(1).replace(",", ""))
|
|
|
|
# Extract final metrics (JSON format from trainer output)
|
|
final_metrics = re.search(r"\{'train_runtime':[^}]+\}", stdout)
|
|
if final_metrics:
|
|
try:
|
|
metrics = eval(final_metrics.group(0)) # Safe: only numbers and strings
|
|
result["final_metrics"] = {
|
|
"train_loss": metrics.get("train_loss"),
|
|
"train_runtime": metrics.get("train_runtime"),
|
|
"train_samples_per_second": metrics.get("train_samples_per_second"),
|
|
}
|
|
except Exception:
|
|
pass
|
|
|
|
# Check completion
|
|
if "Training completed" in stdout:
|
|
result["completed"] = True
|
|
|
|
# 3. Extract warnings (limit to 20)
|
|
warnings = re.findall(r"\[WARNING[^\]]*\][^\n]+", stdout)
|
|
if warnings:
|
|
result["warnings"] = list(set(warnings))[:20]
|
|
|
|
# 4. Fallback: if parsing failed, include truncated raw log
|
|
if not result.get("error") and not result.get("training_started"):
|
|
result["raw_log_tail"] = stdout[-2000:] if len(stdout) > 2000 else stdout
|
|
|
|
return json.dumps(result, indent=2, ensure_ascii=False)
|
|
|
|
def _run_micro_batch_test(self, config_yaml: str, workspace: FBWorkspace, env) -> ValidationResult:
|
|
"""Run micro-batch training test for runtime validation"""
|
|
result = ValidationResult(success=True, filtered_config=config_yaml)
|
|
ws_prefix = get_workspace_prefix(env)
|
|
|
|
# Create micro-batch test configuration
|
|
config = yaml.safe_load(config_yaml)
|
|
if not isinstance(config, dict):
|
|
result.success = False
|
|
result.execution_output = "Invalid YAML configuration"
|
|
result.errors.append("Invalid configuration for micro-batch test")
|
|
return result
|
|
|
|
test_config = config.copy()
|
|
|
|
# Load extra test parameters from workspace (generated by coder in 2nd turn)
|
|
extra_test_params = yaml.safe_load(workspace.file_dict[FT_TEST_PARAMS_FILE_NAME])
|
|
|
|
# Merge extra test parameters (overrides previous settings)
|
|
if extra_test_params:
|
|
test_config.update(extra_test_params)
|
|
|
|
# Run micro-batch training
|
|
workspace.inject_files(**{FT_DEBUG_YAML_FILE_NAME: yaml.dump(test_config, default_flow_style=False)})
|
|
training_result = workspace.run(
|
|
env=env,
|
|
entry=f"llamafactory-cli train {FT_DEBUG_YAML_FILE_NAME}",
|
|
)
|
|
|
|
# Remove micro-batch test files
|
|
workspace.remove_files([FT_DEBUG_YAML_FILE_NAME, FT_TEST_PARAMS_FILE_NAME])
|
|
|
|
# Parse and store structured execution output (reduces ~36k tokens to ~500)
|
|
raw_stdout = training_result.stdout if training_result.stdout else ""
|
|
result.raw_stdout = raw_stdout # Keep full log for UI
|
|
result.execution_output = self._parse_execution_log(raw_stdout, training_result.exit_code)
|
|
|
|
# Check results
|
|
progress_indicators = ["train_loss", "Training:", "Epoch", "loss:", "step"]
|
|
has_progress = any(ind.lower() in training_result.stdout.lower() for ind in progress_indicators)
|
|
|
|
if training_result.exit_code == 0 and has_progress:
|
|
logger.info("Micro-batch test passed")
|
|
result.success = True
|
|
else:
|
|
result.success = False
|
|
result.errors.append(f"Micro-batch test failed (exit_code={training_result.exit_code})")
|
|
|
|
return result
|