Unusually high initial loss during continual pre-training of the Gemma2-2B model.
State: closed | Pull request: no
Labels: name: stale

Could someone please share the details on how to properly set up the Gemma2 pre-training? Many thanks in advance! :)

The SFT scripts are performing well (training loss < 3), but the pre-training loss is extremely high (training loss > 20).

### Experiments
The following experiments are running with the same docker `nvcr.io/nvidia/nemo:24.05.llama3.1` and the same dataset `databricks-dolly-15k`.

![image](https://github.com/user-attachments/assets/94d4e204-dd74-426e-adf5-4636b7513d20)

**SFT training:**
Following the [official document](https://docs.nvidia.com/nemo-framework/user-guide/latest/llms/gemma/sft.html), I set up the SFT training.
[A] gemma2_2b_sft_bos_rc2:
-codebase: https://github.com/NVIDIA/NeMo/tree/r2.0.0rc1
-data: databricks-dolly-15k/train.jsonl
-add bos: Yes

[B]gemma2_2b_sft_0830_main:  
-codebase: https://github.com/NVIDIA/NeMo/tree/1ce9089143b0136523cb08bb37941a35c9b08307
-data: databricks-dolly-15k/train.jsonl
-add bos: Yes

**Continual pre-training**
I didn't find the pre-training recipe for gemma. Thus I just preprocess the dataset with following cmd.
```
python ${NEMO_CODE_PATH}/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \
    --input ${DATASET_PATH}/sample.jsonl \
    --json-keys text \
    --tokenizer-library huggingface \
    --dataset-impl mmap \
    --tokenizer-type ${MODEL_PATH}/gemma2_2b \
    --output-prefix ${DATASET_PATH}/sample_gemma2_preprocessed/sample \
    --append-eod \
    --workers=1
```
[C]gemma2_2b_sft_pretraining
-codebase: https://github.com/NVIDIA/NeMo/tree/1ce9089143b0136523cb08bb37941a35c9b08307
-data: databricks-dolly-15k/train.jsonl
-add bos: No

[D]gemma2_2b_sft_pretraining_bos
-codebase: https://github.com/NVIDIA/NeMo/tree/1ce9089143b0136523cb08bb37941a35c9b08307
-data: databricks-dolly-15k/train.jsonl
-add bos: Yes
I modify [preprocess_data_for_megatron.py](https://github.com/NVIDIA/NeMo/blob/1ce9089143b0136523cb08bb37941a35c9b08307/scripts/nlp_language_modeling/preprocess_data_for_megatron.py#L178) as following to append the BOS token.
```
  for sentence in Encoder.splitter.tokenize(text):
      sentence_ids = Encoder.tokenizer.text_to_ids(sentence)
      sentence_ids = [Encoder.tokenizer.bos_id] + sentence_ids
      if len(sentence_ids) > 0:
          doc_ids.append(sentence_ids)
  if len(doc_ids) > 0 and self.args.append_eod:
      doc_ids[-1].append(Encoder.tokenizer.eos_id)
  ids[key] = doc_ids
```

### Observations
(1) SFT tuning is generally stable, but varies depending on changes in the codebase. [A] and [B] use the same data and Docker environment, yet they show different loss curves.
(2) plays a crucial role, as highlighted by @Emperorizzis in [comments](https://github.com/NVIDIA/NeMo/pull/9672#issuecomment-2225972135). This was further explored in detail by  [https://unsloth.ai/blog/gemma-bugs](https://unsloth.ai/blog/gemma-bugs).


### Looking for help
Could the official maintainers provide a stable version tag for the codebase and detailed guidance on how to run NeMo for pre-training the Gemma2 model? 
Thank you for your outstanding contributions to the amazing NeMo framework!

_Originally posted by @longxudou in https://github.com/NVIDIA/NeMo/issues/9672#issuecomment-2324228927_
            

This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.

This issue was closed because it has been inactive for 7 days since being marked as stale.