corpus.blog
41554 blogs · 4,417,066 posts Claim your blog
Record · corpus.blog/posts/01a08760-2f45-725c-acef-451bdf18f4a5

JAX-LM: Language Modelling and Distributed Training in JAX

chuyishang.com · published 24 March 2026

A record is what was published and who pointed at it. The text of the post is not held here: read it at the source, or at its Wayback capture.

Read the post at chuyishang.com Wayback capture
0blogs citing
First cited
—
Most recent
—
Rank this month
not ranked
Links from this post
22
Words captured
18,007

Cited by

No blog we hold has cited this post yet.

Links from this post

Link
Host
https://cs336.stanford.edu/
cs336.stanford.edu
https://jax-ml.github.io/scaling-book/
jax-ml.github.io
https://github.com/chuyishang/jax-lm/tree/assignment
github.com
https://github.com/chuyishang/jax-lm/tree/main/jax_impl/distributed
github.com
JAX-LM Assignment Read at the source ↗
chuyishang.com
https://github.com/chuyishang/jax-lm/tree/main/jax_impl
github.com
https://github.com/chuyishang/jax-lm/tree/main/pytorch_impl
github.com
https://github.com/chuyishang/jax-lm/tree/main/pytorch_impl/distributed
github.com
https://github.com/stanford-cs336/assignment1-basics
github.com
https://docs.jax.dev/en/latest/index.html
docs.jax.dev
https://flax.readthedocs.io/en/latest/index.html
flax.readthedocs.io
https://openxla.org/xla
openxla.org
https://docs.jax.dev/en/latest/jaxpr.html
docs.jax.dev
https://pytorch.org/blog/overview-of-pytorch-autograd-engine/
pytorch.org
https://flax.readthedocs.io/en/stable/why.html
flax.readthedocs.io
https://github.com/chuyishang/jax-lm/blob/main/jax_impl/basic/model.py
github.com
https://github.com/chuyishang/jax-lm/blob/main/jax_impl/basic/train.py
github.com
https://jax-ml.github.io/scaling-book/sharding/
jax-ml.github.io
https://cloud.google.com/blog/products/ai-machine-learning/faster-distributed-training-with-google-clouds-reduction-server
cloud.google.com
https://jax-ml.github.io/scaling-book/training/
jax-ml.github.io
https://chuyishang.com/assets/img/jax-lm/fsdp+tp.jpg
chuyishang.com
https://github.com/chuyishang/jax-lm/
github.com