Skip to navigation
Training

Automated Hyperparameter Optimisation — LLM Tuning + Grid Search

View as Markdown

Hyperparameter Tuning Guidance

When starting with a new dataset, tune parameters in this order — earlier parameters have larger impact and adjusting them makes later parameters more stable.

Step 1 — Validate basic learning: Set epochs: 1, num_boost_round: 50. The goal is to confirm the pipeline runs end to end and produces a non-random PR-AUC (> 0.5). If AUC stays at 0.5, check data quality: label distribution, feature normalisation, index correctness.

Step 2 — Tune GNN epochs and learning rate: Increase epochs until validation AUC plateaus. If AUC drops after plateau, you are overfitting — stop earlier. Learning rates between 0.001 and 0.01 work well for most datasets.

Step 3 — Tune XGBoost num_boost_round and learning_rate: A smaller learning rate (0.01–0.05) with more rounds (500–1000) usually outperforms a large rate with fewer rounds. Set log_period: 5 to see per-round progress clearly. Important: when using a slow learning rate, increase early_stopping_rounds proportionally (e.g. early_stopping_rounds: 200) — at LR=0.01 the per-round improvement is very small and the default patience of 50 rounds may trigger early stopping before the model converges.

Step 4 — Tune XGBoost regularisation for overfitting: If training PR-AUC is much higher than validation PR-AUC, increase min_child_weight (try 5, 10, 20), set gamma to 0.1–1.0, and reduce subsample to 0.7–0.8.

Step 5 — Try a deeper GNN or different encoder: Once the baseline is stable, experiment with num_gnn_layers: 3 (and a corresponding third value in num_neighbors), or switch to encoder: gat.

Both tools ship inside the container as /workspace/mnmg/llm_tune.sh and /workspace/mnmg/grid_search.py. No extra software is needed on the host — you only need Docker, your data, and a config file.

Strategy Overview

PhaseToolWhen to useTypical duration
Phase 1 — LLM tuningllm_tune.shFresh dataset; explore the space quickly3–5 training runs
Phase 2 — Grid searchgrid_search.pyRefine around the best config from Phase 110–50 targeted runs

Run both phases sequentially. Phase 1 narrows you down to a good neighbourhood; Phase 2 exhaustively tests the combinations inside that neighbourhood.

Get the Image

Both tools are baked into the image. Pull it from NGC:

docker pull nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0

Phase 1 — LLM-Guided Iterative Tuning

llm_tune.sh runs N iterations of a train → analyse → propose → apply loop entirely inside the container. Each iteration:

  1. Trains the model with the current working config
  2. Saves a per-iteration snapshot: metrics_iter_NN.json
  3. Tracks the best config by validation PR-AUC — saves it to best_llm_tuned_config.yaml whenever a new best is found
  4. Queries an LLM with the parsed training log — the LLM reads loss curves, grad norms, XGBoost eval history, embedding statistics, and validation metrics
  5. Shows the proposed config changes with a reason for each change
  6. Asks for approval (y / N / q) — or applies automatically in --auto mode
  7. Repeats from step 1 with the improved config

At the end, two config files are written to the output directory:

FileDescription
best_llm_tuned_config.yamlConfig from the iteration with the highest validation PR-AUC — use this for Phase 2
final_llm_tuned_config.yamlConfig from the last iteration (most tuned, not necessarily best)

A metrics comparison table is printed at the end showing val PR-AUC, test PR-AUC, and F1 across all iterations.

The mounted config.yaml is never modified — the script copies it to a local working path inside the container at startup. You may mount it with or without :ro.

LLM Provider Selection

The script accepts any OpenAI-compatible endpoint. Pass credentials as environment variables. All three LLM_* variables must be set together to activate the custom endpoint; if any is missing the script falls back to the --provider registry (defaults to Anthropic). If an API key is invalid the script retries the next available endpoint automatically.

ScenarioVariables to set
Anthropic Claude (default)ANTHROPIC_API_KEY
OpenAILLM_API_KEY + LLM_MODEL + LLM_BASE_URL
Any OpenAI-compatible APILLM_API_KEY + LLM_MODEL + LLM_BASE_URL

Automated Mode — NVIDIA NIM (Recommended)

docker run --rm -i --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  --entrypoint bash \
  -e LLM_API_KEY=<your-api-key> \
  -e LLM_BASE_URL=<endpoint-url> \
  -e LLM_MODEL=<model-id> \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/llm_tune.sh --iterations 5 --auto
--entrypoint bash must come before the image name — it is a docker run flag, not a command argument. The script /workspace/mnmg/llm_tune.sh is passed as the command after the image name.

Substitute the three env vars for your provider:

ProviderLLM_API_KEYLLM_BASE_URLLLM_MODEL
NVIDIA NIMnvapi-...https://inference-api.nvidia.comnvidia/meta/llama-3.3-70b-instruct
OpenAIsk-...https://api.openai.com/v1gpt-4o

Anthropic: Use ANTHROPIC_API_KEY=sk-ant-... instead of LLM_API_KEY — Anthropic’s API is not OpenAI-compatible and is accessed using the --provider anthropic path, not the LLM_* env vars.

Interactive Mode (Review Each Suggestion Before Applying)

Drop --auto and replace -i with -it to review and approve each LLM suggestion:

docker run --rm -it --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  --entrypoint bash \
  -e LLM_API_KEY=<your-api-key> \
  -e LLM_BASE_URL=<endpoint-url> \
  -e LLM_MODEL=<model-id> \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/llm_tune.sh --iterations 5

Custom Tuning Prompt

The LLM system prompt is baked into the container at /workspace/mnmg/tune_prompt.txt. The full content is in Appendix C.

To override it without rebuilding the image, extract the default first and edit it:

# Step 1: extract the default prompt from the container
docker run --rm --entrypoint cat nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/tune_prompt.txt > my_prompt.txt
# Step 2: edit my_prompt.txt to suit your dataset / objectives
# Step 3: mount your custom prompt when running llm_tune.sh
docker run --rm -i --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  --entrypoint bash \
  -e LLM_API_KEY=<your-api-key> \
  -e LLM_BASE_URL=<endpoint-url> \
  -e LLM_MODEL=<model-id> \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  -v /path/to/my_prompt.txt:/workspace/mnmg/tune_prompt.txt:ro \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/llm_tune.sh --iterations 5 --auto

llm_tune.sh Options Reference

OptionDefaultDescription
-n, --iterations N5Number of train → tune cycles
-c, --config PATHauto-detectConfig file path inside the container
-o, --output-dir DIRfrom configTraining output directory
--autooffApply all suggestions without interactive review
--nproc Nfrom configGPU processes per node
--provider NAMEanthropicProvider when LLM_* vars are not set
--model MODELprovider defaultOverride model within the provider

What the Output Looks Like

══════════════════════════════════════════════════════════════════════
[llm_tune] Iteration 1/5  —  2025-05-01 10:00:00
══════════════════════════════════════════════════════════════════════
[llm_tune] Starting training …
[llm_tune] Training complete in 847s.
[llm_tune] New best  iter=1  val_pr_auc=0.881  → best_llm_tuned_config.yaml
[llm_tune] Training log: /workspace/output/training_20250501_100000.log
[llm_tune] Querying LLM for hyperparameter suggestions …
## Proposed Changes
- gnn.learning_rate: 0.005 → 0.010 — loss still decreasing at last epoch, need faster convergence
- gnn.epochs: 3 → 5 — still_decreasing=true, model not converged
- xgb.num_boost_round: 256 → 384 — reached_max_rounds=true, more rounds needed
[llm_tune] Auto mode — applying changes.
...
══════════════════════════════════════════════════════════════════════
[llm_tune] LLM tuning loop complete (5 iteration(s)).
══════════════════════════════════════════════════════════════════════
[llm_tune] Final tuned config → /workspace/output/final_llm_tuned_config.yaml
[llm_tune] Best config        → /workspace/output/best_llm_tuned_config.yaml
[llm_tune]                      iter=3  val_pr_auc=0.921
 Iter  Val PR-AUC  Test PR-AUC  Val Logloss  F1 (best thr)
────────────────────────────────────────────────────────────
   01       0.881        0.791            -          0.813
   02       0.903        0.821            -          0.841
   03       0.921        0.835            -          0.857
   04       0.908        0.819            -          0.843
   05       0.895        0.810            -          0.831

Use best_llm_tuned_config.yaml (iteration 3 above) as the base for Phase 2 grid search — not final_llm_tuned_config.yaml.

Phase 2 — Grid Search

grid_search.py runs a Cartesian product of GNN and XGBoost parameter values. Each trial writes its own config to /tmp/grid_search/trial_NNN/config.yaml inside the container and saves training output to output/grid_search/trial_NNN/. At the end — and after any interruption — a ranked leaderboard is written to output/grid_search/leaderboard.json and leaderboard.csv.

The config file is read-only for grid search — mount it with :ro.

With the Built-In Default Grid

The default grid tests a narrow range around a typical starting config:

GNN:  hidden_channels ∈ {64, 128}  ×  learning_rate ∈ {0.001, 0.005, 0.01}  ×  epochs ∈ {8, 16}  ×  num_gnn_layers ∈ {1, 2, 3}  ×  encoder ∈ {sage, gat}
XGB:  max_depth ∈ {4, 6}  ×  learning_rate ∈ {0.1, 0.2}  ×  min_child_weight ∈ {3, 5}  ×  subsample ∈ {0.8, 1.0}  ×  num_parallel_tree ∈ {1, 3}
docker run --rm -it --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  --entrypoint python3 \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/grid_search.py

With a Custom Grid File

Create grid.yaml on your host — list only the parameters you want to vary:

# grid.yaml — narrow grid around best LLM-suggested config
gnn:
  hidden_channels: [128, 256]
  learning_rate:   [0.005, 0.008, 0.010]
  epochs:          [5, 7]
xgb:
  max_depth:        [5, 6, 7]
  learning_rate:    [0.08, 0.10, 0.12]
  min_child_weight: [1, 2]
docker run --rm -it --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  -v /path/to/grid.yaml:/workspace/grid.yaml:ro \
  --entrypoint python3 \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/grid_search.py --grid /workspace/grid.yaml

Limit Trial Count With Random Sampling

When the Cartesian product is too large, --max-trials randomly samples a subset:

docker run --rm -it --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  -v /path/to/grid.yaml:/workspace/grid.yaml:ro \
  --entrypoint python3 \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/grid_search.py \
    --grid /workspace/grid.yaml \
    --max-trials 20 \
    --random-seed 42

Resume an Interrupted Run

--resume works at the grid level, not per-trial. When passed, grid_search.py checks every trial’s output directory for a log containing "Training complete!". Trials that have it are skipped; all others are re-run from scratch.

Example — grid interrupted at trial 7 of 20:

trial_001  ✓ complete  → skipped, metrics collected from existing log
trial_002  ✓ complete  → skipped
...
trial_006  ✓ complete  → skipped
trial_007  ✗ incomplete → re-runs from epoch 1
trial_008  ✗ missing   → runs
...
trial_020  ✗ missing   → runs
--resume does not resume a trial mid-epoch. A trial that was killed halfway through GNN training restarts from epoch 1. To avoid repeating the GNN phase, set skip_gnn_train: true in the base config.yaml before resuming (only safe if a completed GNN checkpoint exists in checkpoint_dir or output_dir).
docker run --rm -it --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  -v /path/to/grid.yaml:/workspace/grid.yaml:ro \
  --entrypoint python3 \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/grid_search.py \
    --grid /workspace/grid.yaml \
    --max-trials 20 \
    --resume

Re-Run a Specific Trial

Each trial saves its full merged config (base config + trial overrides) to output/grid_search/trial_NNN/training_config.<ext> where <ext> matches the extension of the base config (.yaml, .yml, or .json). Use it to re-run that exact trial independently:

docker run --rm --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/output/grid_search/trial_002:/workspace/output \
  -v /path/to/output/grid_search/trial_002/training_config.yaml:/workspace/config.yaml:ro \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0

This is useful for verifying the best trial with a longer run or different random seed without re-running the full grid.

Dry-Run — Preview the Trial Plan Without Training

docker run --rm --gpus all \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/grid.yaml:/workspace/grid.yaml:ro \
  --entrypoint python3 \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/grid_search.py \
    --grid /workspace/grid.yaml \
    --max-trials 20 \
    --dry-run

grid_search.py Options Reference

OptionDefaultDescription
--config PATHauto-detectBase config file (YAML or JSON)
--grid PATHbuilt-inGrid definition YAML
--output-dir DIR{config.output_dir}/grid_searchResults root
--max-trials NunlimitedRandom sample when total exceeds N
--random-seed N42Seed for --max-trials sampling
--nproc Nfrom configGPU processes per node
--dry-runoffPrint trial plan without running
--resumeoffSkip already-completed trials

Searchable Parameters

GNN (gnn section in grid YAML): hidden_channels, num_gnn_layers, encoder, heads, concat, epochs, batch_size, learning_rate

num_neighbors is auto-generated from num_gnn_layers when not listed explicitly — each additional layer halves the previous neighbour count.

XGBoost (xgb section in grid YAML): num_boost_round, num_parallel_tree, max_depth, learning_rate, subsample, colsample_bytree, min_child_weight, gamma

Output Structure

/path/to/output/
└── grid_search/
    ├── leaderboard.json        ← all trials ranked by val_pr_auc (then val_f1)
    ├── leaderboard.csv         ← same data, spreadsheet-friendly
    ├── trial_001/
    │   ├── training_*.log      ← full training log
    │   ├── torchrun.log        ← raw process stdout/stderr
    │   ├── xgboost_fraud.json  ← trained XGBoost model
    │   └── metrics.json        ← validation metrics for this trial
    ├── trial_002/
    │   └── ...
    └── trial_NNN/
        └── ...

Leaderboard Sample

  Grid Search Leaderboard
  ══════════════════════════════════════════════════════════════════
 rank │ trial │ status │ val_pr_auc │ val_f1 │ gnn.hidden_channels │ ...
─────┼───────┼────────┼────────────┼────────┼─────────────────────┼───
 1   │ 007   │ ok     │ 0.9641     │ 0.8823 │ 256                 │ ...
 2   │ 003   │ ok     │ 0.9588     │ 0.8741 │ 128                 │ ...
 3   │ 011   │ ok     │ 0.9521     │ 0.8655 │ 256                 │ ...

Copy the hyperparameters from the top-ranked row into your config.yaml for production training.

# ── Phase 1: LLM tuning — find a good region quickly ──────────────────────────
#
# The script copies config.yaml to an internal working copy at startup.
# The original mounted file is never modified.
# Tuned configs are written to the output directory (see table above).
# Remove --auto to review each proposal interactively.
docker run --rm --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml \
  -v /path/to/output:/workspace/output \
  -e ANTHROPIC_API_KEY=sk-ant-... \
  --entrypoint /workspace/mnmg/llm_tune.sh \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  --iterations 5 --auto
# Tuned configs are in /path/to/output/: best_llm_tuned_config.yaml + final_llm_tuned_config.yaml
# The original config.yaml mount is never modified.
# ── Phase 2: Grid search — squeeze out the last few AUCPR points ───────────────
#
# Mount config.yaml as read-only — grid search derives per-trial configs from it.
# Mount your custom grid.yaml to test values around the LLM-suggested settings.
docker run --rm --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  -v /path/to/grid.yaml:/workspace/grid.yaml:ro \
  --entrypoint python3 \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/grid_search.py \
    --grid /workspace/grid.yaml \
    --max-trials 20
# Results are in /path/to/output/grid_search/leaderboard.json and leaderboard.csv.
# Copy the top-ranked trial's hyperparameters into config.yaml, then run a final
# full training with the standard entrypoint:
docker run --rm --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /path/to/data:/data:ro \
  -v /path/to/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/output:/workspace/output \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0

Running with Your Dataset — Single Node

This section gives ready-to-run commands. Replace the three host paths shown below with your actual locations — everything else is handled automatically.

Mount pointWhat to put hereContainer path
Your graph data directory/path/to/your/data/data (read-only)
Your config file/path/to/your/config.yaml/workspace/config.yaml
Output directory/path/to/your/output/workspace/output

Step 0 — Verify GPU Count

Check how many GPUs are visible to the container before you start:

docker run --rm --gpus all nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 nvidia-smi --list-gpus

Then set num_gpus in your config.yaml to match. Both Phase 1 and Phase 2 read this value automatically — no flags needed.

# config.yaml — set this to the number of GPUs on your node
models:
  - kind: GNN_XGBoost
    num_gpus: 2        # ← change to 1, 4, 8, etc.

Step 1 — Phase 1: LLM Tuning on Your Data

The script copies the mounted config to an internal working path at startup and never writes back to the mount. You may mount config with or without :ro. At the end of the loop, two configs are written to the output directory: best_llm_tuned_config.yaml (highest val PR-AUC) and final_llm_tuned_config.yaml (last iteration).

docker run --rm -it --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  --entrypoint bash \
  -e LLM_API_KEY=<your-api-key> \
  -e LLM_BASE_URL=<endpoint-url> \
  -e LLM_MODEL=<model-id> \
  -v /path/to/your/data:/data:ro \
  -v /path/to/your/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/your/output:/workspace/output \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/llm_tune.sh --iterations 5

Each iteration prints the proposed changes and asks [y/N/q]. Add --auto (and replace -it with -i) to apply all suggestions without prompts:

docker run --rm -i --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  --entrypoint bash \
  -e LLM_API_KEY=<your-api-key> \
  -e LLM_BASE_URL=<endpoint-url> \
  -e LLM_MODEL=<model-id> \
  -v /path/to/your/data:/data:ro \
  -v /path/to/your/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/your/output:/workspace/output \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/llm_tune.sh --iterations 5 --auto

When the container exits, the tuned configs are in the output directory: /path/to/your/output/best_llm_tuned_config.yaml (best by val PR-AUC) and /path/to/your/output/final_llm_tuned_config.yaml (last iteration). The original /path/to/your/config.yaml is unchanged.

Step 2 — Phase 2: Grid Search on Your Data

Create grid.yaml on your host with values to test around whatever Phase 1 settled on:

# grid.yaml — narrow grid around the best LLM-suggested config
gnn:
  hidden_channels: [64, 128, 256]
  learning_rate:   [0.003, 0.005, 0.010]
  epochs:          [3, 5]
xgb:
  max_depth:        [4, 6, 8]
  learning_rate:    [0.05, 0.10, 0.20]
  min_child_weight: [1, 3]

Config is mounted with :ro — grid search writes per-trial configs internally and never modifies the base config.

docker run --rm -it --gpus all \
  --shm-size=10g \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -v /path/to/your/data:/data:ro \
  -v /path/to/your/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/your/output:/workspace/output \
  -v /path/to/your/grid.yaml:/workspace/grid.yaml:ro \
  --entrypoint python3 \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0 \
  /workspace/mnmg/grid_search.py \
    --grid /workspace/grid.yaml \
    --max-trials 20

Results land in /path/to/your/output/grid_search/:

leaderboard.json   ← all trials ranked by val_pr_auc
leaderboard.csv    ← same, spreadsheet-friendly
trial_001/         ← per-trial logs and model artifacts
trial_002/
...

Step 3 — Final Production Training With the Best Config

Copy the top-ranked trial’s hyperparameters into config.yaml, then run the standard entrypoint:

docker run --rm --gpus all \
  --shm-size=10g \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -v /path/to/your/data:/data:ro \
  -v /path/to/your/config.yaml:/workspace/config.yaml:ro \
  -v /path/to/your/output:/workspace/output \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0