Out of memory

[Yurri-hub]

I use 2 RTX3090 for block-nerf dataset training, and I find that it runs out of memory in moe module. Is that normal?

Althought I decrease the batch_size, even to batch_size=2312, it still went out of memory.
And my scripts are like:

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch \
    --use_env --master_port=12345 --nproc_per_node=2 -m \
    switch_nerf.train \
    --config=switch_nerf/configs/switch_nerf/mission_bay.yaml \
    --use_moe --exp_name=data/outputs \
    --dataset_path=data/datasets/waymo-block-nerf/v1.0 \
    --block_train_list_path=switch_nerf/datasets/lists/block_nerf_train_val.txt \
    --block_image_hash_id_map_path=switch_nerf/datasets/lists/block_nerf_id_map.json \
    --chunk_paths=data/chunks/mission_bay_chunk_radii_1 \
    --no_bg_nerf --near=0.01 --far=10.0 \
    --use_balance_loss \
    --i_print=200 \
    --batch_size=12312 \
    --moe_expert_type=expertmlp \
    --moe_train_batch \
    --moe_test_batch \
    --model_chunk_size=212992 \
    --coarse_samples=257 \
    --fine_samples=257 \
    --moe_capacity_factor=1.0 \
    --batch_prioritized_routing \
    --moe_l_aux_wt=0.0005 \
    --amp_use_bfloat16 \
    --use_moe_external_gate \
    --use_gate_input_norm \
    --use_sigma_noise \
    --sigma_noise_std=1.0

[MiZhenxing]

Hi, batch size of 12312 for 2 3090 is very large. I typically use 13312 for 8 GPUs, i.e. 1664 for each GPU. You may need to decrease the batch size.