High Memory Consumption (~80GB) with Batch Size=1 on SOC Dataset: Is this expected?

[WendyYW]

Hello,

I am currently using this project to train on the SOC (Spin-Orbit Coupling) dataset. I have encountered a significant memory consumption issue where the VRAM usage is unexpectedly high.

Current Behavior I am running the training on a node equipped with 8x NVIDIA H100 (80GB) GPUs. Even with batch_size=1, the GPU memory usage almost hits the limit, reaching approximately 80,305 MiB per GPU.

Environment:

Hardware: NVIDIA H100 80GB HBM3 (x8)

CUDA Version: 12.4

Dataset: SOC Data

Screenshots

Questions

Is this memory usage expected? Given the complexity of the SOC data, is it normal for a single sample to consume ~80GB of VRAM?

Config Check: Could you please check my attached config? Are there any specific parameters (e.g., related to layer size or precision) that I should adjust to lower the memory footprint?

Optimization: Any suggestions on how to optimize this to allow for a larger batch size?

Thank you for your help!

Configuration Below is the configuration file I am using. I suspect there might be some parameters here contributing to the high memory usage.

dataset_params:
  batch_size: 1
  split_file: null
  test_ratio: 0.1
  train_ratio: 0.8
  val_ratio: 0.1
  graph_data_path: ./test_dir
  num_workers: 16
losses_metrics:
  losses:
  - loss_weight: 27.211
    metric: mae
    prediction: hamiltonian
    target: hamiltonian
  metrics:
  - metric: mae
    prediction: hamiltonian
    target: hamiltonian
optim_params:
  lr: 0.01
  lr_decay: 0.5
  lr_patience: 5
  gradient_clip_val: 0.0
  max_epochs: 6000
  min_epochs: 100
  stop_patience: 30

output_nets:
  output_module: HamGNN_out
  HamGNN_out:
    ham_only: true # true: Only the Hamiltonian H is computed; 'false': Fit both H and S
    ham_type: openmx # openmx: fit openmx Hamiltonian; abacus: fit abacus Hamiltonian
    nao_max: 26 # The maximum number of atomic orbitals in the data set, which can be 14, 19 or 27
    add_H0: True # Generally true, the complete Hamiltonian is predicted as the sum of H_scf plus H_nonscf (H0)
    symmetrize: True # if set to true, the Hermitian symmetry constraint is imposed on the Hamiltonian
    calculate_band_energy: False # Whether to calculate the energy bands to train the model
    num_k: 5 # When calculating the energy bands, the number of K points to use
    band_num_control: 8 # `dict`: controls how many orbitals are considered for each atom in energy bands; `int`: [vbm-num, vbm+num]; `null`: all bands
    k_path: null # `auto`: Automatically determine the k-point path; `null`: random k-point path; `list`: list of k-point paths provided by the user
    soc_switch: True # if true, fit the SOC Hamiltonian
    nonlinearity_type: gate # norm or gate
    # spin constrained
    spin_constrained: False
    collinear_spin: False
    minMagneticMoment: 0.5

profiler_params:
  progress_bar_refresh_rat: 1
  train_dir: ./train_output_1w #The folder for saving training information and prediction results. This directory can be read by tensorboard to monitor the training process.

representation_nets:
  # Network parameters usually do not need to be changed.
  HamGNN_pre:
    cutoff: 26.0
    cutoff_func: cos
    edge_sh_normalization: component
    edge_sh_normalize: true
    ######## Irreps set 1 (crystal): ################
    irreps_edge_sh: 0e + 1o + 2e + 3o + 4e + 5o
    irreps_node_features: 64x0e+64x0o+32x1o+16x1e+12x2o+25x2e+18x3o+9x3e+4x4o+9x4e+4x5o+4x5e+2x6e
    num_layers: 3
    num_radial: 64
    num_types: 96
    rbf_func: bessel
    set_features: true
    radial_MLP: [64, 64]
    use_corr_prod: False
    correlation: 2
    num_hidden_features: 16
    use_kan: False
    radius_scale: 1.01
    build_internal_graph: False

setup:
  GNN_Net: HamGNNpre
  accelerator: null
  ignore_warnings: true
  checkpoint_path: null # Path to the model weights file
  load_from_checkpoint: False
  resume: False
  num_gpus: 8 # null: use cpu; [i]: use the ith GPU device
  precision: 32
  property: hamiltonian
  stage: fit # fit: training; test: inference

[QuantumLab-ZY]

Dear @WendyYW,

Could you share how many atoms are included in each structure in your training set?
If the number of atoms per structure is below ~50, then the error you’re seeing is unlikely to be caused by insufficient GPU memory. Instead, a more probable cause is host CPU memory exhaustion when running DDP training with 8 GPUs.

Older versions of HamGNN load the entire graph_data.npz file into memory. When using 8 processes for DDP, this results in 8 copies of the full dataset in memory, which can easily exceed the system’s RAM capacity.

If memory usage is indeed the issue, I recommend switching to HamGNN 2.1 and converting your dataset from graph_data.npz to the graph_data.lmdb format. The LMDB format supports streaming reads, significantly reducing memory consumption during multi-GPU training.

Below is a script you can use to convert graph_data.npz to graph_data.lmdb:

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Script to convert graph_data.npz format to LMDB format
Usage: python npz_to_lmdb.py --input graph_data.npz --output graph_data.lmdb
"""
import numpy as np
import lmdb
import pickle
import os
import argparse
import time
from tqdm import tqdm
import shutil
def convert_npz_to_lmdb(npz_path, lmdb_path, map_size=None):
    """
    Convert NPZ file to LMDB format
    
    Parameters:
    npz_path: str, path to input NPZ file
    lmdb_path: str, path to output LMDB database
    map_size: int, memory size allocated for LMDB database (in bytes), default is 1TB or automatically estimated
    """
    # Ensure output directory exists
    os.makedirs(os.path.dirname(os.path.abspath(lmdb_path)), exist_ok=True)
    
    # Load original data
    print(f"Loading NPZ file: {npz_path}")
    start_time = time.time()
    
    try:
        data = np.load(npz_path, allow_pickle=True)
        graphs = data['graph'].item()  # Get dictionary object
    except Exception as e:
        print(f"Cannot load NPZ file: {e}")
        return False
    
    print(f"NPZ loading complete, time elapsed: {time.time() - start_time:.2f} seconds")
    print(f"Found {len(graphs)} graph data")
    
    # Estimate required map_size
    if map_size is None:
        # First sample to calculate average size
        sample_size = min(100, len(graphs))
        sample_indices = np.random.choice(list(graphs.keys()), sample_size, replace=False)
        
        total_bytes = 0
        for idx in sample_indices:
            # Serialize object and calculate size
            serialized = pickle.dumps(graphs[idx])
            total_bytes += len(serialized)
        
        # Estimate total size (average size * number of graphs * safety factor)
        avg_size = total_bytes / sample_size
        estimated_size = int(avg_size * len(graphs) * 1.5)  # Add 50% safety margin
        
        # Ensure at least 1GB, not exceeding 1TB
        map_size = max(1_073_741_824, min(estimated_size, 1_099_511_627_776))
        print(f"Estimated database size: {map_size / (1024**3):.2f} GB")
    
    # If target path already exists, create a backup
    if os.path.exists(lmdb_path):
        backup_path = lmdb_path + ".backup"
        print(f"Target LMDB path already exists, creating backup: {backup_path}")
        if os.path.exists(backup_path):
            shutil.rmtree(backup_path)
        shutil.move(lmdb_path, backup_path)
    
    # Create LMDB environment
    env = lmdb.open(lmdb_path, map_size=map_size)
    
    # Write data
    print("Starting data conversion to LMDB...")
    start_time = time.time()
    
    with env.begin(write=True) as txn:
        # Store the number of graphs
        txn.put('num_graphs'.encode(), str(len(graphs)).encode())
        
        # Create progress bar using tqdm
        for i, idx in enumerate(tqdm(sorted(graphs.keys()), desc="Conversion progress")):
            graph = graphs[idx]
            
            # Serialize graph data
            serialized_data = pickle.dumps(graph)
            
            # Store to LMDB
            txn.put(f'graph_{i}'.encode(), serialized_data)
    
    # Close environment
    env.close()
    
    # Output results
    print(f"Conversion completed! Time elapsed: {time.time() - start_time:.2f} seconds")
    print(f"LMDB database has been saved to: {lmdb_path}")
    
    # Validate data
    validate_lmdb(lmdb_path)
    
    return True
def validate_lmdb(lmdb_path):
    """Validate that the LMDB database is correctly created"""
    print("Validating LMDB database...")
    
    try:
        env = lmdb.open(lmdb_path, readonly=True, lock=False)
        with env.begin() as txn:
            num_graphs = int(txn.get('num_graphs'.encode()).decode())
            print(f"LMDB contains {num_graphs} graph data")
            
            # Try to read the first and last data points
            first_key = 'graph_0'.encode()
            if txn.get(first_key) is not None:
                print("✅ Successfully read the first data item")
            
            last_key = f'graph_{num_graphs-1}'.encode()
            if txn.get(last_key) is not None:
                print("✅ Successfully read the last data item")
        
        env.close()
        print("✅ LMDB validation successful!")
    except Exception as e:
        print(f"❌ LMDB validation failed: {e}")
def main():
    parser = argparse.ArgumentParser(description='Convert NPZ graph data to LMDB format')
    parser.add_argument('--input', type=str, required=True, help='Input NPZ file path')
    parser.add_argument('--output', type=str, required=True, help='Output LMDB database path')
    parser.add_argument('--map-size', type=int, default=None, help='Memory size allocated for LMDB database (bytes)')
    args = parser.parse_args()
    
    # Check if input file exists
    if not os.path.exists(args.input):
        print(f"Error: Input file does not exist: {args.input}")
        return 1
    
    # Convert data
    success = convert_npz_to_lmdb(args.input, args.output, args.map_size)
    
    return 0 if success else 1
if __name__ == '__main__':
    exit(main())

Please let me know if the issue persists after switching to the LMDB format—I’m happy to help further.

Best,
Yang Zhong

---