[DLRM v2] How to modify the default training script of DLRM v2 to train the model with limited GPU memory

[JJingL]

Hi Teams,
I have run the default training script of DLRM v2 to train the model, however, the GPU I used doesn't have enough memory for the default setting. I just modified the training script with the following changes:
--num_embeddings_per_feature 26000000,39060,17295,7424,20265,3,7122,1543,63,26000000,3067956,405282,10,2209,11938,155,4,976,14,26000000,26000000,26000000,590152,12973,108,36 \

However，the output of eval_accuracy didn't increase, and the final result is around 0.70x. Is there anyone have any idea?

Ps：Here is the exact command I've tried:
export TOTAL_TRAINING_SAMPLES=4195197692
export GLOBAL_BATCH_SIZE=16384
export WORLD_SIZE=8

torchx run -s local_cwd dist.ddp -j 1x8 --script dlrm_main.py --
--embedding_dim 128
--dense_arch_layer_sizes 512,256,128
--over_arch_layer_sizes 1024,1024,512,256,1
--in_memory_binary_criteo_path /workspace/DLRM/numpy_contiguous_shuffled_output_dataset_dir
--num_embeddings_per_feature 24000000,39060,17295,7424,20265,3,7122,1543,63,24000000,3067956,405282,10,2209,11938,155,4,976,14,24000000,24000000,24000000,590152,12973,108,36
--validation_freq_within_epoch $((TOTAL_TRAINING_SAMPLES / (GLOBAL_BATCH_SIZE * 40)))
--epochs 1
--adagrad
--pin_memory
--mmap_mode
--batch_size $((GLOBAL_BATCH_SIZE / WORLD_SIZE))
--interaction_type=dcn
--dcn_num_layers=3
--dcn_low_rank_dim=512
--learning_rate 0.004
--shuffle_batches
--multi_hot_distribution_type uniform
--multi_hot_sizes=3,2,1,2,6,1,1,1,1,7,3,8,1,6,9,5,1,1,1,12,100,27,10,3,1,1
--print_sharding_plan

[menghongtao]

Hi I am also faced the problem same with you, I changed the embeddingdim ,and the result auc of eval is 0.70*, have you solved your issued