Add Deepseek-v3

[miladm]

Goal:

• Add Deepseek-v3
• Enable single-chip TPU functionality
• Use TorchAx

Non-Goal / Next-Steps:

• Real input tensor
• Real weights
• FP8 quantization kernels enablement (fp8_gmm)
• Distrbuted

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[miladm]

Should probably submit the original model before this PR to easily spot the diff.

[qihqi]

Great to get it running! Few high level changes:

1. let's remove all the README.md etc along with their perf graphs, pdfs etc. Instead, in the .py files that we forked, add the URL to their github repo of where we forked from.
2. Make a new README.md with the command you run to run it on TPU
3. edit requirements.txt to TPU requirements.

[qihqi]

stamp to unblock feel free to merge after the change

[tengyifei]

Can we add a unit test to make sure it runs and produces correct results as compared to CPU eager? Example for Llama: https://github.com/AI-Hypercomputer/torchprime/blob/main/torchprime/experimental/torchax_models/test/test_llama.py#L40

Btw, we also need to fix lint. You can do that by running ruff check and ruff format.

Thanks!

[yaochengji]

@miladm looks like only part of the model is converted to jax device.

I reverted the code change in model.py and tried to call model.to("jax") outside. I got the error: torchax.tensor.OperatorNotFound: Operator with name aten::rms_norm has no lowering.

I can rewrite the rms_norm op to fine grained op as llama model first. @qihqi , BTW, do we apply decomposition before converting torch op to jax ops?