Training | NVIDIA Financial Fraud NIM Documentation

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 Your Own Config (Recommended)

Mount your graph data, output directory, and config file:

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 \
  -v /your/config.yaml:/workspace/config.yaml:ro \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0

The container automatically:

Detects your config at /workspace/config.yaml.
Validates it before any GPU work starts — configuration errors print a clean message and exit immediately.
Reads num_gpus from the config and launches torchrun with the correct number of workers.
Runs the full pipeline and writes all artifacts to /workspace/output.

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.2 Config-File Mode (Recommended)

Config files produce reproducible, self-documenting runs that are easy to version-control and share with teammates. The pipeline validates the config before launching any GPU processes, so configuration errors surface immediately with a clean error message.

Config Fields (with example values)

The block below shows every supported config field with inline descriptions. Use it as a reference when building your own config file.

paths:
  data_dir: /data                   # path inside the container where your graph data is accessible;
                                    # must match the container-side of your -v mount:
                                    #   -v /your/host/data:/data:ro   # ← /data is the container side
  output_dir: /workspace/output     # path inside the container where all training artifacts are written
                                    # (model weights, XGBoost model, Triton repo, metrics);
                                    # mount a host directory here so output survives after the container exits:
                                    #   -v /your/host/output:/workspace/output
                                    # when save_node_embeddings: true, node embeddings are also
                                    # saved here under node_embeddings/
  partition_dir: /tmp/fraud_partitions  # PyArrow partition cache; for multi-node Slurm runs this
                                    # needs to be bind-mounted from a shared filesystem (e.g. Lustre) dir 
  embedding_dir: null               # XGBoost edge-embedding temp dir (memmap: true); auto if null
  checkpoint_dir: null              # path to model_final.pt for warm-start; null = train from scratch
  skip_partition: false             # set true to reuse existing partitions (same dataset only)
models:
  - kind: GNN_XGBoost               # GNN_XGBoost | GNN_XGBoost_NP | XGBoost
    gpu: multi                      # single | multi
    num_gpus: 2                     # GPUs per node; required when gpu: multi; ignored under Slurm
    hyperparameters:
      gnn:
        hidden_channels: 128        # hidden layer width; larger = more expressive, slower
        num_gnn_layers: 2           # message-passing hops; must equal len(num_neighbors)
        num_neighbors: [25, 10]     # neighbours sampled per hop (outermost first)
        encoder: sage               # sage | gat | transformer | general
        heads: 4                    # attention heads; gat/transformer only
        concat: true                # true = concat heads; false = average; gat/transformer only
        dropout: 0.0                # dropout between GNN layers; 0 = disabled
        epochs: 5                   # training epochs
        batch_size: 1024            # seed edges per mini-batch
        learning_rate: 0.005        # GNN optimizer learning rate
        focal_gamma: 0.0            # focal loss exponent; 0 = standard BCE; >0 down-weights easy negatives
        focal_alpha: null           # fraud-class weight in (0,1); null = disabled
        no_raw_node_features_in_embedding: false   # true = GNN embeddings only (no raw features)
        only_raw_node_features_in_embedding: false  # true = raw features only (no GNN embeddings)
        inf_num_neighbors: null     # neighbour limits during embedding extraction; defaults to num_neighbors
                                    #   set lower to cap latency on high-degree nodes
        inf_batch_size: null        # mini-batch size during GNN embedding extraction; default batch_size // 8
        skip_gnn_train: false       # true = skip training, load from checkpoint_dir, go to XGBoost
        # ── Saved node embeddings (EP hybrid inference) ──────────
        save_node_embeddings: false # true = run Stage 9: save per-node embeddings to node_embeddings/
                                    #   enables client_production.py / client_saved_emb.py fast path
        final_fit_epochs: 2         # fine-tune epochs on val+test before export; 0 = skip
        final_fit_lr_scale: 0.1     # LR multiplier for fine-tuning (lr × scale)
        emb_dtype: float32          # float32 | float16 (halves disk footprint)
        merge_chunk_size: null      # row-chunk when merging shards; null = full shard in RAM
                                    #   set to e.g. 1000000 to cap RAM for large graphs
      xgb:
        num_boost_round: 256        # number of boosting rounds
        max_depth: 6                # max tree depth; deeper = more interactions, higher overfit risk
        learning_rate: 0.1          # XGBoost shrinkage; lower = more rounds needed
        subsample: 1.0              # row sub-sampling ratio per tree
        colsample_bytree: 1.0       # column sub-sampling ratio per tree
        min_child_weight: 1         # min instance weight in a child (regularisation)
        gamma: 0.0                  # min loss reduction for a split; 0 = no constraint
        num_parallel_tree: 1        # trees per round; >1 = random-forest-style boosting
        base_score: 0.5             # initial prediction score; set to TRAINING fraud rate
        scale_pos_weight: null      # positive-class weight; null = auto (1-rate)/rate
        prior_test: null            # production fraud rate; when set, val set is downsampled to
                                    #   simulate this rate and the threshold is calibrated on it
        eval_metric: [auc, aucpr, logloss]  # last metric drives early stopping
        early_stopping_rounds: 50   # stop if no improvement over last N rounds; null = disabled
        # ── Memory management (large graphs) ─────────────────────
        batched: false              # true = stream embeddings in batches (needed when matrix > CPU RAM)
        extmem: false               # true = ExtMemQuantileDMatrix; requires batched: true, XGBoost >= 3.0
        batch_size: 65536           # edges per GPU batch when batched: true
        memmap: false               # true = write GNN embeddings to disk; requires batched: true
        cache_host_ratio: null      # fraction of extmem data pinned in host RAM; null = auto
        log_period: 10              # print eval metrics every N rounds

YAML Config (recommended format)

The example below shows a realistic production config using the GAT encoder with saved node embeddings and batched XGBoost memory mode enabled:

paths:
  data_dir: /data                   # path inside the container where your graph data is accessible;
                                    # must match the container-side of your -v mount:
                                    #   -v /your/host/data:/data:ro   # ← /data is the container side
  output_dir: /workspace/output     # path inside the container where all training artifacts are written
                                    # (model weights, XGBoost model, Triton repo, metrics);
                                    # mount a host directory here so output survives after the container exits:
                                    #   -v /your/host/output:/workspace/output
                                    # when save_node_embeddings: true, node embeddings are also
                                    # saved here under node_embeddings/
  partition_dir: /tmp/fraud_partitions  # PyArrow partition cache; for multi-node Slurm runs this
                                    # needs to bind-mounted from the shared  filesystem (e.g. Lustre) dir
  checkpoint_dir: null              # set to resume from a previous run's model_final.pt
  skip_partition: false             # set true to reuse existing partitions
models:
  - kind: GNN_XGBoost
    gpu: multi
    num_gpus: 2                     # match your hardware; ignored under Slurm (--gpus-per-node wins)
    hyperparameters:
      gnn:
        hidden_channels: 128        # embedding width per node; 64–256 covers most use cases
        num_gnn_layers: 2           # aggregation hops (must equal len(num_neighbors))
        num_neighbors: [16, 8]      # neighbours sampled per hop (outermost first)
        encoder: gat                # sage | gat | transformer | general
        heads: 4                    # attention heads (gat/transformer only)
        concat: true                # true = concat heads; false = average
        dropout: 0.3                # dropout between GNN layers (0 = disabled)
        epochs: 2                   # training epochs; more epochs help with class imbalance
        batch_size: 4096            # seed edges per mini-batch
        learning_rate: 0.05
        focal_gamma: 0.0            # >0 enables focal loss (e.g. 1.0–3.0); 0 = standard BCE
        focal_alpha: null           # fraud-class weight in (0,1); null = disabled
        no_raw_node_features_in_embedding: false
        only_raw_node_features_in_embedding: false
        skip_gnn_train: false
        # ── Saved node embeddings (EP_EMB hybrid inference) ──────
        save_node_embeddings: true  # true = run; enables fast memmap
                                    # inference path (client_production.py / client_saved_emb.py)
        final_fit_lr_scale: 0.2     # LR multiplier for the final-fit phase (lr × scale)
        final_fit_epochs: 4         # fine-tune epochs on val+test before embedding export
      xgb:
        num_boost_round: 512
        num_parallel_tree: 3        # random-forest-style variance reduction per round
        max_depth: 6
        learning_rate: 0.05
        subsample: 0.9
        colsample_bytree: 0.95
        min_child_weight: 8
        gamma: 0.05
        # ── Class imbalance calibration ───────────────────────────
        base_score: 0.5             # ← set to training fraud rate (e.g. 0.09 for 9% fraud)
        scale_pos_weight: 4         # null = auto-computed from training data (recommended)
        prior_test: 0.002           # ← set to production fraud rate for calibrated scores
        eval_metric: [auc, aucpr, logloss]
        early_stopping_rounds: 50
        # ── Memory management (for large graphs) ─────────────────
        batched: true               # stream embeddings in batches; recommended for large graphs
        extmem: true                # XGBoost external memory mode (requires XGBoost >= 3.0)
        batch_size: 65536
        memmap: true                # write GNN embeddings to disk instead of RAM
        cache_host_ratio: 0.5       # fraction of extmem data pinned in host RAM
        log_period: 10

JSON Config

{
  "paths": {
    "data_dir": "/data",
    "output_dir": "/workspace/output",
    "partition_dir": "/tmp/fraud_partitions",
    "embedding_dir": "/tmp/fraud_emb_dir",
    "checkpoint_dir": null,
    "skip_partition": false
  },
  "models": [
    {
      "kind": "GNN_XGBoost",
      "gpu": "multi",
      "num_gpus": 2,
      "hyperparameters": {
        "gnn": {
          "hidden_channels": 128,
          "num_gnn_layers": 2,
          "num_neighbors": [8, 4],
          "encoder": "sage",
          "heads": 4,
          "concat": true,
          "epochs": 5,
          "batch_size": 1024,
          "learning_rate": 0.005
        },
        "xgb": {
          "num_boost_round": 256,
          "max_depth": 6,
          "learning_rate": 0.1,
          "subsample": 1.0,
          "colsample_bytree": 1.0,
          "min_child_weight": 1,
          "log_period": 10
        }
      }
    }
  ]
}

XGBoost-Only Mode (No GNN)

Use kind: XGBoost to train gradient-boosted trees directly on node feature vectors, with no graph context. This is useful as a performance baseline to quantify the value added by the GNN.

paths:
  data_dir: /data
  output_dir: /workspace/output
models:
  - kind: XGBoost
    gpu: multi
    num_gpus: 2
    data:
      label_column: fraud      # column name in your tabular CSV label file
      format: csv              # csv | parquet | orc
    hyperparameters:
      num_boost_round: 300
      max_depth: 6
      learning_rate: 0.1
      subsample: 0.8
      colsample_bytree: 0.8
      min_child_weight: 2
      gamma: 0.1
      log_period: 10

Running with a Config File

Mount your config to /workspace/config.yaml (YAML or JSON — both accepted):

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 \
  -v /path/to/configs/production.yaml:/workspace/config.yaml:ro \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0

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.