About OOM During Training and Questions Regarding Attn

[Zheng-Jay]

Thank you for your contribution! I have encountered some issues.
1、Full train
Here is my training script:

CUDA_VISIBLE_DEVICES="0,5" torchrun --nproc_per_node 2 \
-m training.run \
--output_dir ./output/7-2_full \
--model_name_or_path /mnt/data1/zmj/embedding_model/GritLM-7B/model \
--train_data /mnt/data1/zmj/embedding_model/gritlm-main/gritlm/training/toy_data/test_7-2.jsonl \
--learning_rate 1e-5 \
--num_train_epochs 1 \
--per_device_train_batch_size 1 \
--dataloader_drop_last True \
--normalized True \
--temperature 0.02 \
--query_max_len 32 \
--passage_max_len 128 \
--train_group_size 2 \
--mode embedding \
--attn cccc

Why do I get an OOM (Out of Memory) error? My GPU is 80G A800, and the model is only 7B with a batch size of 1. I believe this configuration should not cause an OOM.

2、LoRA train
To be able to perform training, I used the --lora option. However, after training, the checkpoint saved is 24GB, while the original model was only 14GB:

27G     /mnt/data1/zmj/embedding_model/gritlm-main/gritlm/output/7-2_lora/checkpoint-800

I would like to know why this is the case. Additionally, I received the following warning when loading:
Some weights of the model checkpoint at /mnt/data1/zmj/embedding_model/gritlm-main/gritlm/output/7-2_lora were not used when initializing MistralForCausalLM: ['model.base_model.model.embed_tokens.weight', 'model.base_model.model.layers.0.input_layernorm.weight', 'model.base_model.model.layers.0.mlp.down_proj.weight', 'model.base_model.model.layers.0.mlp.gate_proj.weight',...]

3、attn
After reading the paper, I understand that you used bidirectional attn for training the embedding task. However, why does the example script you provided for the embedding task use: --attn cccc

I look forward to your response.

[Muennighoff]

1,
try adding the below to save memory

    --attn_implementation sdpa \
    --no_gen_gas \
    --no_emb_gas \
    --split_emb \

2,
I have not tried LoRA but it looks to me like your checkpoint was saved in FP32 which doubles the size. The warning is problematic cuz it means ur weights are not loaded.

3,

it is just an example script; the actual training script is here: https://github.com/ContextualAI/gritlm/blob/main/scripts/training/train_gritlm_7b.sh

[Zheng-Jay]

1, try adding the below to save memory1, 尝试添加以下内容以节省内存

    --attn_implementation sdpa \
    --no_gen_gas \
    --no_emb_gas \
    --split_emb \

2, I have not tried LoRA but it looks to me like your checkpoint was saved in FP32 which doubles the size. The warning is problematic cuz it means ur weights are not loaded.2, 我没有尝试过 LoRA，但在我看来，您的检查点保存在 FP32 中，大小增加了一倍。警告是有问题的，因为它意味着您的砝码没有加载。

3,3,

it is just an example script; the actual training script is here: https://github.com/ContextualAI/gritlm/blob/main/scripts/training/train_gritlm_7b.sh它只是一个示例脚本;实际的训练脚本在这里：https://github.com/ContextualAI/gritlm/blob/main/scripts/training/train_gritlm_7b.sh

Thank you very much for your reply. I'm not very familiar with training code, but I will give it a try. Thanks again!