NNX Migration for Mistral models

[NuojCheng]

Description

Migrate Mistral model to NNX. Comparing f4f286c (before mistral migration) and 5984d9c (after mistral migration).

Tests

Training tests

Command

python -m MaxText.train MaxText/configs/base.yml \
    run_name=nc_test_mistral_$RANDOM \
    steps=5 \
    base_output_directory=gs://chengnuojin-maxtext-logs \
    dataset_path=gs://chengnuojin-maxtext-dataset \
    model_name=mistral-7b \
    enable_checkpointing=False \
    per_device_batch_size=1

Webdiff

https://diff.googleplex.com/#key=qvnNP8tANEIK

Decode Inference tests

Command

python3 -m MaxText.decode src/MaxText/configs/base.yml \
model_name=mistral-7b \
tokenizer_path=assets/tokenizer.mistral-v1 \
tokenizer_type=sentencepiece \
scan_layers=false \
per_device_batch_size=1 \
ici_fsdp_parallelism=1 \
ici_autoregressive_parallelism=-1 \
max_prefill_predict_length=128 \
max_target_length=256 \
prompt="I love to" \
attention=dot_product \
load_parameters_path=gs://chengnuojin-maxtext-logs/chengnuojin_decode_32458/checkpoints/0/items

Webdiff

https://diff.googleplex.com/#key=dbqiUSGbkqov

Jetstream Inference test

Command

Step 1: https://paste.googleplex.com/6590290169298944

Step 2: https://paste.googleplex.com/6674235103772672

Step 3: https://paste.googleplex.com/6448284356968448

Before Mistral Migration

xprof

Memstats: After load_params:
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_0(process=0,(0,0,0,0))
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_1(process=0,(1,0,0,0))
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_2(process=0,(0,1,0,0))
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_3(process=0,(1,1,0,0))

RAMstats: After load_params:
        Using (GB) 27.17 / 440.83 (6.163374%) -->  Available:410.84

After Mistral Migration

xprof

Memstats: After load_params:
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_0(process=0,(0,0,0,0))
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_1(process=0,(1,0,0,0))
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_2(process=0,(0,1,0,0))
        Using (GB) 3.37 / 95.74 (3.519950%) on TPU_3(process=0,(1,1,0,0))

RAMstats: After load_params:
        Using (GB) 27.94 / 440.83 (6.338044%) -->  Available:410.06

Web diff

https://diff.googleplex.com/#key=siw4m0hLYUAT

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[bvandermoon]

LGTM just a couple small comments. Are the inference decode test results in the description from after the rebase with #2370? I see some very slight memory diff

Also discussed this offline already, but do you mind running some Maxengine/Jetstream test from a checkpoint?

[RissyRan]

I see both of accuracy is 0 from JetStream. Did you kill the process or others?

Results

{'accuracy': 0.0, 'gen_num': 5000}

[NuojCheng]

I see both of accuracy is 0 from JetStream. Did you kill the process or others?

Results

{'accuracy': 0.0, 'gen_num': 5000}

I did not kill the process. I am not sure what happened here, maybe using the wrong tokenizer? (I used the mistral tokenizer under assets instead of huggingface in llama2)

[bvandermoon]

LGTM after the testing you have been continuing. Just one comment

[bvandermoon]

LGTM, thanks for the great thorough testing so far. Please do follow up here when you have the results of the golden logits test you mentioned you were running