pin for future tokenizer

Merge branch 'main' of github.com:LargeWorldModel/lwm into main

updates

raise an exception of max_sequence_length is smaller than the input sequence

Fix inference

use scan to control remat

update ringattention version

simplify code with importing ringattention

switch to autotokenizer for simplicity

updated jax version

Fix ring attention inference GPU issue

Add data and sharding docs

Added PyTorch sampling script + more detailed descriptions for running scripts

Add needle haystack data script

Merge pull request #26 from ruslanmustafin/main