Advanced Training Options#

XGBoost Memory Modes#

By default, all edge embeddings are collected in CPU RAM before XGBoost training. For large datasets this can exceed available memory. The following modes progressively reduce memory pressure.

Mode	Typical use case	Config	Memory usage
Default (in-memory)	Embeddings fit in CPU RAM	(no special flags)	Full dataset in RAM
Batched QuantileDMatrix	RAM is tight, any XGBoost version	`batched: true`	One batch in RAM at a time during histogram building
Batched ExtMemQuantileDMatrix	RAM is very tight, XGBoost ≥ 3.0	`batched: true, extmem: true`	XGBoost manages its own disk cache
Batched + Memmap	Embeddings cannot fit in RAM even as batches	`batched: true, memmap: true`	Embeddings on disk; one batch read at a time
Memmap + Extmem	Maximum memory savings	`batched: true, memmap: true, extmem: true`	Minimal RAM; both embedding store and XGBoost cache on disk

Choosing batch_size for batched modes:

batch_size × embedding_dimension × 4 bytes ≤ target GPU memory for data loading

Example:
  embedding_dimension = 256 (2 × 128 hidden + raw features)
  target GPU memory   = 8 GB

  batch_size ≤ 8 GB / (256 × 4 B) = 8,388,608 → use 4,194,304 or lower

A good starting point is batch_size: 65536 (≈ 64 MB for 256-dim embeddings) and double it until you reach the memory limit.

Config example — batched external memory:

xgb:
  num_boost_round: 512
  max_depth: 6
  learning_rate: 0.05
  batched: true
  batch_size: 131072
  extmem: true
  cache_host_ratio: 0.5   # pin 50% of data in host RAM between rounds

How Edge Embeddings Are Generated Before XGBoost Trains#

Understanding when to use memmap and extmem requires knowing how the GNN produces the embedding matrix that XGBoost trains on.

memmap: false (default) — in-RAM accumulation:

The GNN processes training edges in mini-batches. After each batch’s forward pass, the resulting edge embedding tensor is appended to a list held in CPU RAM. When all batches complete, the list is concatenated into a single tensor and handed to XGBoost. No disk I/O occurs. This is fastest but requires enough system RAM to hold the full embedding matrix at once (all training edges × embedding dimension × 4 bytes).

memmap: true — incremental disk write (requires batched: true):

Each mini-batch’s embeddings are written immediately to a disk-backed memory-mapped file (edge_emb_train_rank0.mmap) as the GNN produces them. XGBoost then reads from this file in batches — either using QuantileDMatrix (all data in disk file, but processed batch-by-batch) or using ExtMemQuantileDMatrix when extmem: true is also set (XGBoost manages its own disk cache too). Peak RAM during training equals approximately one batch of embeddings rather than the full training set.

The .mmap file is temporary — it is deleted automatically after XGBoost finishes training.

Config	Peak RAM during XGBoost training	Disk needed during training
`batched: false` (default)	Full embedding matrix in RAM	None
`batched: true, extmem: false`	One batch + XGBoost histogram buffers	Full embedding matrix on disk
`batched: true, extmem: true`	One batch	Full embedding matrix + XGBoost cache on disk

Use memmap: true (requires batched: true) in RAM-constrained environments when system RAM cannot comfortably hold all training edge embeddings at once. extmem: true is a separate flag that enables XGBoost’s own ExtMemQuantileDMatrix disk cache; the two options are independent and can be combined.

Resuming Interrupted Training#

Set checkpoint_dir in your config to warm-start GNN weights from a previous run. The GNN saves a single checkpoint (model_final.pt) to output_dir after all epochs complete. If training is interrupted before finishing, no partial checkpoint is available and training will restart from the beginning unless a previous run’s model_final.pt already exists in output_dir or checkpoint_dir.

paths:
  checkpoint_dir: /workspace/checkpoints/run1

To retrain only the XGBoost step using the saved GNN embedding (without re-running GNN training):

paths:
  checkpoint_dir: /workspace/checkpoints/run1

# in models[0]:
hyperparameters:
  gnn:
    skip_gnn_train: true       # skip GNN training; load checkpoint and re-run embedding extraction for XGBoost

Monitoring Training Progress#

Real-Time Console Output#

All log messages from rank 0 appear on stdout. Key lines to watch:

INFO — Config schema validation passed ✓
INFO — world_size=2  kind=GNN_XGBoost
INFO — Partitioning graph across 2 GPUs ...
INFO — Epoch 01, Loss: 0.4821  lr=0.005000  grad_norm=0.8432  time=12.3s
INFO — Epoch 02, Loss: 0.3409  lr=0.005000  grad_norm=0.7201  time=11.9s
...
INFO — [0]	train-auc:0.91234	train-aucpr:0.88765	train-logloss:0.18234	val-auc:0.87654	val-aucpr:0.84543	val-logloss:0.22876
INFO — [9]	train-auc:0.99123	train-aucpr:0.98765	train-logloss:0.01234	val-auc:0.94321	val-aucpr:0.93210	val-logloss:0.09876
...
INFO — Training complete!

Log Files#

Each run writes a timestamped log file: output_dir/training_YYYYMMDD_HHMMSS.log

# Watch live
tail -f /workspace/output/training_*.log

# Search for errors
grep -n "ERROR\|Traceback\|FAILED" /workspace/output/training_*.log

Metrics File#

After training completes, output_dir/metrics.json contains the full evaluation report:

{
  "pr_auc": 0.8998,
  "best_threshold": 0.7651,
  "best_threshold_raw": 0.7651,
  "calibrated": false,
  "threshold_0.5": {
    "accuracy": 0.9543,
    "precision": 0.6952,
    "recall": 0.8653,
    "f1": 0.7710
  },
  "threshold_best": {
    "accuracy": 0.9718,
    "precision": 0.9013,
    "recall": 0.7665,
    "f1": 0.8285
  },
  "test": {
    "pr_auc": 0.9243,
    "val_threshold": 0.7651,
    "threshold_val": {
      "accuracy": 0.9757,
      "precision": 0.8310,
      "recall": 0.8788,
      "f1": 0.8542
    },
    "threshold_0.5": {
      "accuracy": 0.9612,
      "precision": 0.7241,
      "recall": 0.8910,
      "f1": 0.7989
    }
  }
}

When prior_test is configured, the file gains six additional top-level fields:

{
  "prior_train": 0.10,
  "prior_test": 0.0003,
  "val_sim_positives": 3,
  "val_sim_total": 10000,
  "threshold_0.5_prod": {
    "accuracy": 0.9992,
    "precision": 0.4231,
    "recall": 0.8412,
    "f1": 0.5631
  },
  "threshold_best_prod": {
    "accuracy": 0.9996,
    "precision": 0.6841,
    "recall": 0.7109,
    "f1": 0.6972
  }
}

threshold_best_prod is the F1-optimal threshold evaluated on a class-ratio-simulated validation set (see Class Imbalance and Calibration). Use this value — not best_threshold — as the production decision threshold when your live fraud rate differs significantly from the training fraud rate.

Metric glossary:

Metric	Meaning	Target
`pr_auc`	Area under the Precision-Recall curve. Summarises model quality across all thresholds. Insensitive to class imbalance.	> 0.85 is good; > 0.90 is excellent
`best_threshold`	Decision threshold (on the fraud probability) that maximises F1 on the validation set. This value is embedded into the Triton deployment.	Between 0.3 and 0.9 is typical
`precision`	Of all transactions flagged as fraud, what fraction were actually fraudulent?	Depends on business tolerance for false positives
`recall`	Of all truly fraudulent transactions, what fraction did the model catch?	Higher is better; typical production target ≥ 0.80
`f1`	Harmonic mean of precision and recall. Useful single-number summary.	> 0.80 is good for fraud

Class Imbalance and Calibration#

Production fraud rates (0.01%–0.2%) are typically 50–500× lower than training fraud rates (10–20%). If left unaddressed this gap produces overconfident probability scores and a misaligned decision threshold. Three parameters work together to close it:

`scale_pos_weight` — Training Balance#

Upweights the fraud class during XGBoost training to compensate for imbalance in the training data. The default null computes it automatically as (1 − fraud_rate) / fraud_rate from the training labels. Override only if you want a custom ratio — for example to trade off precision vs. recall:

xgb:
  scale_pos_weight: null   # auto (recommended for most cases)
  # scale_pos_weight: 200  # manual: 200:1 non-fraud:fraud

`base_score` — XGBoost’s Initial Prediction#

Sets XGBoost’s intercept before any trees are built. The default 0.5 is calibrated for balanced datasets. Set it to your training fraud rate so the trees only need to learn the residual signal:

xgb:
  base_score: 0.001   # 0.1% training fraud rate

This reduces the number of trees needed for convergence and improves probability calibration on the training distribution.

`prior_test` — Production Threshold Calibration#

When your evaluation environment (OOT test set, live production) has a much lower fraud rate than training, set prior_test to calibrate scores for that environment. The pipeline downsamples validation-set positives so their rate matches prior_test, then finds the F1-optimal threshold on that simulated distribution — only the threshold is recalibrated. AUC is unaffected.

xgb:
  base_score: 0.10       # training fraud rate
  prior_test: 0.0003     # 0.03% expected production fraud rate

This is essential for meaningful precision/recall at production thresholds — without it, the decision threshold is calibrated on a 10%-fraud validation set and will be misaligned when the model is deployed against 0.03% production traffic, resulting in far too many false positives or false negatives depending on where the threshold falls.

Combined Example — Large Imbalance Gap#

models:
  - kind: GNN_XGBoost
    hyperparameters:
      gnn:
        focal_gamma: 1.0    # focus GNN loss on hard fraud examples
        focal_alpha: 0.25   # upweight fraud class in focal loss
      xgb:
        base_score: 0.10    # training data: 10% fraud
        scale_pos_weight: null  # auto = 9x fraud-class weight (for 10% training fraud rate)
        prior_test: 0.0003  # production: 0.03% fraud → recalibrate threshold

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.

Saving Node Embeddings for Hybrid Inference (EP Only)#

Setting gnn.save_node_embeddings: true enables a third inference mode — hybrid inference — that pre-computes GNN node embeddings at training time and uses fast memmap lookup at inference time for any node that was in the training graph.

What Happens at the End of Training#

After the main GNN+XGBoost training loop completes, two additional phases run automatically:

Phase A — Fine-tuning (optional): If final_fit_epochs > 0, the saved model_final.pt is reloaded and trained for final_fit_epochs more epochs on the combined train+val+test edge set, using learning_rate × final_fit_lr_scale as the LR. This improves embedding quality by exposing the model to the full labeled dataset before export.

Phase B — Stage 9: Node embedding export: A full GNN forward pass runs over all nodes of each type (both source and destination node types). Per-rank shard files are written:

output/node_embeddings/
├── {src_type}/
│   ├── rank=0_ids.pt     ← global node IDs for this rank's node partition
│   ├── rank=0_emb.pt     ← GNN output vectors for this rank's node partition
│   ├── rank=0_feat.pt    ← raw node feature vectors (same partition)
│   ├── rank=1_ids.pt
│   ├── rank=1_emb.pt
│   ├── rank=1_feat.pt
│   ├── embeddings.mmap   ← merged flat memmap (created at first client startup)
│   ├── features.mmap     ← merged raw feature memmap (created at first client startup)
│   └── shape.json        ← n_total, hidden_channels, dtype
├── {dst_type}/
│   └── ...
└── manifest.json         ← n_raw_features, world_size, emb_dtype per node type

Note on merge timing: Stage 8 (Triton export) runs in train_gnn_xgb.py before Stage 9 (node embedding export). Because node_embeddings/ does not yet exist when Stage 8 runs, the auto-merge inside triton_export.py is always skipped during a fresh training run. The merge into embeddings.mmap and features.mmap happens lazily at first client startup (load_mmaps() in client_saved_emb.py / client_production.py calls merge_all() automatically when only rank shard .pt files are present). If you need the merged files before starting a client (for example, when copying artifacts to a serving machine), run python infer/merge_embeddings.py --emb_dir <output>/node_embeddings manually. After the merge, the rank shard .pt files are no longer needed at inference time.

Inference Routing#

At inference time, each edge is routed based on whether its src and dst node IDs fall within the training graph:

Condition	Path	Mechanism
Both `src_id` and `dst_id` ≤ training graph bounds	Known — fast path	`np.memmap` lookup → XGBoost FIL directly. No GNN computation.
Either endpoint outside training bounds	Novel — GNN fallback	Full GNN through Triton `fraud_pipeline` ensemble.

The per-path metrics in client_saved_emb.py output tell you what fraction of your inference traffic is hitting each path and the quality of each path independently.

When to Use Saved Node Embeddings#

Use save_node_embeddings: true when:

You need low-latency inference for most edges
Your inference traffic is dominated by nodes seen during training (the known path)
You want to decouple the GNN from the serving critical path

Use the standard client.py (no saved embeddings) when:

You want the simplest possible serving setup
Most inference nodes are novel (cold-start rate > 50%)
Latency requirements allow the full GNN inference time