p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/sample.sh (shellscript)
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python sample.py \
    --batch_size 1 \
    --seq_len 2 \
    --start_frame 1 \
    --checkpoint $ckpt_dir \
    --data_dir $data_dir
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

mkdir -p $ckpt_dir

export XLA_FLAGS=--xla_gpu_autotune_level=0
# srun python train_tokenizer.py \
PYTHONUNBUFFERED=1 python train_tokenizer.py \
    --batch_size 12 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] train_tokenizer.py (python)
    # model_state = nnx.state(optimizer.model)
    # model_shardings = nnx.get_named_sharding(model_state, mesh)
    # model_sharded_state = jax.lax.with_sharding_constraint(model_state, model_shardings)
    # nnx.update(optimizer.model, model_sharded_state)
    # optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    # optimizer_shardings = nnx.get_named_sharding(optimizer_state, mesh)
    # optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, optimizer_shardings)
    # nnx.update(optimizer, optimizer_sharded_state)

    model_state = nnx.state(optimizer.model)
    model_sharded_state = jax.lax.with_sharding_constraint(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)

    jax.debug.visualize_array_sharding(optimizer.model.encoder.blocks[0].attn.ffn_dense1.kernel.value)
[content] train_tokenizer.py (python)
    
    # Debug: Print what shardings we're getting
    print("Model shardings from nnx.get_named_sharding:")
    print(model_shardings)
    
    # Debug: Check if we're getting replicated sharding for all parameters
    def print_sharding_info(pytree, prefix=""):
        for key, value in pytree.items():
            if hasattr(value, 'sharding'):
                print(f"{prefix}{key}: {value.sharding}")
            elif isinstance(value, dict):
                print_sharding_info(value, prefix + f"{key}.")
    
    print_sharding_info(model_shardings)
[content] train_lam.py (python)
    model_state = nnx.state(optimizer.model)
    model_sharded_state = jax.lax.with_sharding_constraint(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)
[content] train_dynamics.py (python)
    model_state = nnx.state(optimizer.model)
    model_sharded_state = jax.lax.with_sharding_constraint(model_state, replicated_sharding)
    nnx.update(optimizer.model, model_sharded_state)
    optimizer_state = nnx.state(optimizer, nnx.optimizer.OptState)
    optimizer_sharded_state = jax.lax.with_sharding_constraint(optimizer_state, replicated_sharding)
    nnx.update(optimizer, optimizer_sharded_state)