Is this page helpful?

Training#

5.1 Quick Start — CLI#

The training pipeline is packaged as a self-contained Docker image (nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0). No local Python environment is required — all dependencies run inside the container.

With the Built-In Default Config (Smoke Test)#

Omitting the config mount uses the container’s built-in default (2 GPUs, GNN_XGBoost, standard hyperparameters):

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

Use this to verify that the image and your data directory are wired up correctly before committing to a full production run.

Useful Runtime Overrides#

All container defaults can be overridden with -e VAR=value:

# Change log verbosity
-e LOG_LEVEL=DEBUG

# Override number of GPU workers without editing the config
-e NPROC=1

# Point to a config at a non-standard container path
-e CONFIG_FILE=/workspace/my_config.yaml -v /host/my_config.yaml:/workspace/my_config.yaml:ro

# Enable NVLink P2P on clusters that support it (disabled by default for portability)
-e NCCL_P2P_DISABLE=0

All hyperparameters are documented in Section 5.3.

5.3 All Hyperparameters Reference#

GNN Hyperparameters (hyperparameters.gnn)#

Parameter

Type

Default

Description

hidden_channels

int

128

Hidden dimension per GNN layer. Larger values increase model capacity and GPU memory usage proportionally.

num_gnn_layers

int

2

Number of message-passing layers. Each layer expands the receptive field by one hop. Must equal len(num_neighbors).

num_neighbors

list[int]

[25, 10]

Neighbours sampled per hop during training. The list is ordered outermost-to-innermost: [hop_2, hop_1] for a 2-layer GNN. Reducing these is the most effective way to lower GPU memory.

encoder

string

sage

Convolution operator. See Section 5.4.

heads

int

4

Multi-head attention count. Used only by gat and transformer. Ignored for sage and general.

concat

bool

true

If true, concatenate attention heads (output dim = heads × hidden_channels). If false, average them (output dim = hidden_channels).

epochs

int

5

Training epochs. One epoch = one full pass over all training edges.

batch_size

int

1024

Number of query edges per mini-batch during GNN training.

learning_rate

float

0.005

Initial Adam learning rate.

dropout

float

0.0

Dropout probability applied between GNN conv layers. 0.0 disables dropout. Values of 0.1–0.3 can reduce overfitting on small graphs.

focal_gamma

float

0.0

Focusing exponent for focal loss. 0.0 (default) uses standard BCE (EP) or cross-entropy (NP). Values > 0 down-weight easy (non-fraud) examples by (1 − p)^gamma, concentrating learning on hard misclassifications. Typical range: 0.5–2.0.

focal_alpha

float

null

Fraud-class weight in (0, 1) for focal loss. Scales fraud-class loss by alpha and non-fraud by (1 − alpha). null = no alpha weighting. Use in combination with focal_gamma when fraud is very rare.

skip_gnn_train

bool

false

If true, skip GNN training and immediately run embedding extraction using the existing checkpoint at checkpoint_dir. Useful when iterating on XGBoost hyperparameters without retraining the GNN.

no_raw_node_features_in_embedding

bool

false

Exclude raw node feature vectors from the final edge embedding. The embedding will contain GNN-learned representations (2 × hidden_channels) plus any edge attributes.

only_raw_node_features_in_embedding

bool

false

Use only raw concatenated node features, discarding the GNN output entirely. Dimensionality of the embedding = sum of all node feature counts + edge attributes. Reduces the pipeline to a node-feature-only baseline.

inf_batch_size

int

null

Edges per mini-batch during embedding extraction (the GNN inference pass over all edges). Defaults to batch_size // 8. Reduce if extraction runs out of memory.

inf_num_neighbors

list[int]

null

Per-hop neighbour sampling limits used during embedding extraction and evaluation (not training). Defaults to num_neighbors when null. Set explicitly to cap sampling on high-degree nodes (e.g. merchants with 200K+ edges) that would otherwise stall the inference loader. Must have the same length as num_neighbors.

save_node_embeddings

bool

false

When true: after training completes, Stage 9 runs a full forward pass over all nodes of the source and destination node types of the predicted edge (the two endpoint types only; context node types are not embedded) and saves per-rank GNN embedding shards to {output_dir}/node_embeddings/{type}/rank={r}_ids.pt, rank={r}_emb.pt, and rank={r}_feat.pt. The merged embeddings.mmap and features.mmap files are created at first client startup (see Saving Node Embeddings). Enables the hybrid inference clients (client_saved_emb.py, client_production.py).

final_fit_epochs

int

2

Fine-tune model_final.pt on validation + test edges for this many additional epochs at the start of Stage 9, immediately before embedding extraction. 0 skips fine-tuning and exports the final trained weights directly. Only applies when save_node_embeddings: true.

final_fit_lr_scale

float

0.1

Learning rate multiplier for fine-tuning epochs: effective LR = learning_rate × final_fit_lr_scale. Only applies when save_node_embeddings: true and final_fit_epochs > 0.

node_emb_batch_size

int

null

Batch size for the Stage 9 node embedding inference pass. null falls back to inf_batch_size if set, else batch_size // 8. Only applies when save_node_embeddings: true.

emb_dtype

string

"float32"

Storage dtype for saved node embeddings: "float32" (default) or "float16". float16 halves disk and memory footprint at negligible accuracy cost; the inference clients always upcast to float32 before the XGBoost call. Only applies when save_node_embeddings: true.

node_emb_cache_size

int

null

LRU in-memory cache size (number of node embeddings per type). When set, the most recently accessed N node embeddings are served from Python RAM, bypassing the memmap entirely. null relies on the OS page-cache. Pass --lru_cache_size to the inference clients to activate at serving time.

merge_chunk_size

int

null

Row-chunk size when merging per-rank embedding shards into the flat embeddings.mmap. null loads each rank shard fully before writing. Set to e.g. 1_000_000 to cap peak RAM during merge to approximately chunk_size × hidden_channels × 4 bytes regardless of graph size.

XGBoost Hyperparameters (hyperparameters.xgb)#

Parameter

Type

Default

Description

num_boost_round

int

256

Total number of trees. Higher values improve fit on training data; watch for diminishing returns on validation AUC.

max_depth

int

6

Maximum depth of each tree. Values 4–8 are typical for fraud. Deeper trees overfit more readily.

learning_rate

float

0.1

Step size shrinkage. Lower values (0.01–0.05) require more rounds but often generalise better.

subsample

float

1.0

Row subsampling ratio per round. Values of 0.7–0.9 reduce overfitting and speed training.

colsample_bytree

float

1.0

Column (feature) subsampling ratio per tree. Values of 0.7–0.9 improve robustness on high-dimensional embeddings.

min_child_weight

int

1

Minimum sum of instance weights for a leaf. Increase (e.g., to 5 or 10) if the model overfits the minority fraud class.

gamma

float

0.0

Minimum loss reduction to make a split. Non-zero values (e.g., 0.1–1.0) act as a pruning regularizer.

base_score

float

0.5

XGBoost’s initial prediction score (the prior fraud probability). The default of 0.5 is appropriate for balanced training data. Set this to your training fraud rate (e.g. 0.10 for 10% training fraud) so the trees only need to learn the residual signal. Use prior_test separately for the production rate. See Class Imbalance and Calibration.

scale_pos_weight

float

null

Weight multiplier for the positive (fraud) class. null = auto-computed from training data as (1 − fraud_rate) / fraud_rate. Set explicitly to override the automatic value. Increasing this recall-biases the model; decreasing it precision-biases.

prior_test

float

null

Expected fraud rate in your production/OOT evaluation environment. When set, the pipeline downsamples validation-set positives to match this rate, then finds the F1-optimal threshold on that simulated distribution — only the threshold is recalibrated; AUC is unaffected. Example: set prior_test: 0.0003 when training at 10% fraud but deploying at 0.03%. See Class Imbalance and Calibration.

eval_metric

list[str]

[auc, aucpr, logloss]

XGBoost evaluation metrics logged each round. The last metric controls early stopping — keep logloss last unless you want a different metric to drive the stopping criterion. If omitted from the config, the schema default ["auc", "aucpr", "logloss"] applies; set to [] to disable eval metrics entirely.

early_stopping_rounds

int

50

Stop boosting early if the last metric on the validation set does not improve for N consecutive rounds. Set to null to disable early stopping and always run num_boost_round rounds.

num_parallel_tree

int

1

Trees per round. Set to > 1 to build a random forest at each round instead of a single tree (random forest boosting).

log_period

int

10

Log validation metrics every N rounds.

skip

bool

false

Skip XGBoost training entirely. The GNN embedding extraction still runs; only the classifier step is omitted.

batched

bool

false

Use DataIter-based batched input. Required if embeddings exceed CPU RAM.

batch_size

int

65536

Edges per batch in batched mode. Rule of thumb: batch_size × feat_dim × 4 bytes ≤ available GPU VRAM.

extmem

bool

false

Use ExtMemQuantileDMatrix instead of QuantileDMatrix. Requires batched: true and XGBoost ≥ 3.0.

memmap

bool

false

Stream edge embeddings from a memory-mapped disk file instead of holding them in CPU RAM. Requires batched: true.

cache_host_ratio

float

null

Fraction of external memory data pinned in host RAM between rounds. null lets XGBoost decide. Only applies with extmem: true.

Path Options (paths)#

Parameter

Default

Description

data_dir

required

Root data directory containing nodes/, edges/, and test_gnn/. (Recommended mount point: /data)

output_dir

required

Destination for all outputs: model_final.pt, xgboost_fraud.json, metrics.json, infer/. (Recommended: /workspace/output)

partition_dir

/tmp/fraud_partitions

WholeGraph partition scratch space. Reuse across runs with skip_partition: true.

embedding_dir

auto

GNN embedding staging directory. Auto-generated if null.

checkpoint_dir

null

Path to a directory containing a pre-existing model_final.pt to warm-start GNN weights. Checked as a fallback after output_dir. New checkpoints are always saved to output_dir/model_final.pt after training completes.

skip_partition

false

Skip graph partitioning. Set true after the first run to save time when rerunning with the same data.

5.4 Choosing an Encoder#

Four graph convolution operators are available. Each has different trade-offs in terms of expressiveness, speed, and memory.

Encoder

Algorithm

Heads param

Best for

Notes

sage

GraphSAGE (mean aggregation)

Ignored

Most datasets; production default

Fast, memory-efficient, proven in large-scale fraud detection

gat

Graph Attention Network

Required

Heterogeneous graphs with noisy edges

Learns per-edge attention weights; more expressive but ~2× slower and higher memory

transformer

Graph Transformer

Required

Complex multi-relational graphs

Highest expressiveness; requires the most memory and tuning

general

GeneralConv

Used (no output expansion)

Experimental / ablation studies

Flexible aggregation; heads are used internally but output stays at hidden_channels

Recommendation for production: Start with sage. It trains fastest, uses the least memory, and achieves strong results on the majority of financial fraud datasets. Move to gat only if you have a specific reason to believe that attention weighting will help (e.g., very heterogeneous graphs with many different relationship types and signal quality varying by relation).

Attention head sizing with gat:

With concat: true, intermediate GAT layers expand to heads × hidden_channels internally, but a linear post-projection maps all node embeddings back to hidden_channels before the classifier and XGBoost stage. The downstream XGBoost input dimension is unchanged regardless of heads. GPU memory during training is higher because intermediate activations are heads × hidden_channels wide — consider reducing hidden_channels to 64 when using gat with 4+ heads.