> This page is for version 3.0.0 (Latest) (default).
> For other versions, use one of these documentation indexes:
> - 3.0.0 (Latest) (default): https://docs.nvidia.com/nim/financial-fraud-training/3.0.0/llms.txt
> - 2.0.0: https://docs.nvidia.com/nim/financial-fraud-training/2.0.0/llms.txt
> - 1.0.0: https://docs.nvidia.com/nim/financial-fraud-training/1.0.0/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nim/financial-fraud-training/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nim/financial-fraud-training/_mcp/server.

# 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:**

```yaml
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`.

```yaml
paths:
  checkpoint_dir: /workspace/checkpoints/run1
```

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

```yaml
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`

```bash
# 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:

```json
{
  "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:

```json
{
  "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](#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:

```yaml
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:

```yaml
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.

```yaml
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

```yaml
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.</Note>

### 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