customize vacabulary size to 2560 instead of 2256 to enable fp8 GEMM in CI

[weifengpy]

current status: CI use tokenizer from "./test/assets/test_tiktoken.model", and has vacab size = 2256. FP8 GEMM requires matrix size to be divisible by 16. without any sharding, we have shape 2256 / 16 = 141, and 141 is not divisible by world size 2/4/8

Option 1: customizer or add a new tokenizer in CI with vacab size = 2560 (2560 / 16 = 160). it's enough to test 2-way, 4-way, 8-way TP

Option 2: enable padding for fp8 GEMM at the cost of memeory spike and 20% perf regression

repro
CONFIG_FILE="./train_configs/debug_model.toml" ./run_llama_train.sh and look for the following log
Building llama3 debugmodel with ModelArgs(dim=256, n_layers=8, n_heads=16, n_kv_heads=None, vocab_size=2256

[weifengpy]

cc @tianyu-l for tokenizer ideas

[tianyu-l]

Maybe check #259 which added the test tokenizer. We should be able to generate a new test tokenizer using the script therein, where it seems we can specify vocab_size.
cc: @wanchaol

[awgu]

Why are we doing any fp8 matmul involving vocab_size? Is this for the final output linear?

I think we conventionally do not run that in fp8 for numerics reasons.

[weifengpy]

Why are we doing any fp8 matmul involving vocab_size? Is this for the final output linear?

I think we conventionally do not run that in fp8 for numerics reasons.

good catch! the complain is from model.output. I opened a PR to keep output linear in high precision #469

[weifengpy]

close since this is "fake" requirements - it's common practice to keep model.output in high precision, ie nn.Linear instead of Float8Linear #469