Out of Memory (OOM) During Training a LLaMA 7B Reward Model (8 A800 40GB GPUs)

[qingyiaaaaa]

Description
I am training a LLaMA 7B reward model using 8 A800 GPUs (each with 40GB of memory). Normally, this hardware configuration allows for full fine-tuning of the llama 7b model.While the setup can barely manage one training iteration, the training process consistently runs out of memory (OOM) during validation.

Environment Details
Model: LLaMA 7B Reward Model
GPUs: 8 × A800 (40GB each)
Framework: NeMo
Parallelism Setup: tensor model parallel size = 4, pipeline model parallel size = 2
Precision: 16
Configuration
Here is my script settings:

Problem Analysis
Training can proceed for one iteration but consistently fails with OOM during the validation phase.

Request for Suggestions
Could you provide recommendations or adjustments to settings that might help address this issue?