[Bug] OOM while inferencing on ogbn-papers100M for link prediction

[isratnisa]

Attempting inference on ogbn-papers100M for link prediction is causing an OOM (out-of-memory) issue, as shown in the attached screenshot. As a result, the system is becoming extremely unresponsive. The OOM happens in the evaluation function (val_mrr, test_mrr = self.evaluator.evaluate(None, test_scores, 0)) function here. System can successfully save the node embedding and relational embeddings and exit the program without any issue when the evaluation function is omitted.

Experiment setup:

Dataset: ogbn-papers100M partitioned into 3
Instance: g4dn.metal
Command to run Inferencing:

python3 -u  ~/dgl/tools/launch.py \
        --workspace /graph-storm/inference_scripts/lp_infer \
        --num_trainers 1 \
        --num_servers 1 \
        --num_samplers 0 \
        --part_config /data/ogbn-papers100M-3p/ogbn-papers100M.json \
        --ip_config  /data/ip_list_p3_metal.txt \
        --ssh_port 2222 \
        "python3 lp_infer_gnn.py --cf  /data/ogbn_papers100M_infer_p3.yaml  --use-node-embeddings false --num-gpus 4 --part-config /data/ogbn-papers100M-3p/ogbn-papers100M.json  --restore-model-path /data/papers100M-lp-p3-model/epoch-0  --feat-name feat --no-validation false"

Reproduced with the following environment:

• DGL 1.0.2 + GSF github/gitlab version
• DGL 1.0.0 + GSF github/gitlab version

Smaller dataset like ogbn-mag works fine on similar setup.

[isratnisa]

@classicsong