p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/train.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --job-name=crowd-pilot_qwen3-0.6b_batch_size_1
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# Log the sbatch script
cat $0

cd maxtext
source .venv/bin/activate

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID
BASE_OUTPUT_DIRECTORY=/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/outputs/
PRETRAINED_CHECKPOINT_PATH=/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/checkpoint_pretrained_maxtext/Qwen/Qwen3-0.6B/
GRAIN_TRAIN_FILES="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/array_record/train_*.array_record"
GRAIN_EVAL_FILES="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/array_record/val_*.array_record"

env | grep SLURM

# FIXME(f.srambical): change the batch size to the largest one that can fit in accelerator memory
python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  load_parameters_path=$PRETRAINED_CHECKPOINT_PATH/0/items/ \
  model_name=qwen3-0.6b \
  max_target_length=8192 \
  eval_interval=500 \
  steps=30000 \
  per_device_batch_size=4 \
  tokenizer_path=src/MaxText/assets/qwen3-tokenizer \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  skip_jax_distributed_system=True
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/decode.sh (shellscript)
python3 -m MaxText.decode src/MaxText/configs/base.yml\
    load_parameters_path=gs://your-gcs-bucket/qwen3_maxtext_ckpt/0/items\
    tokenizer_type=huggingface\
    tokenizer_path=src/MaxText/assets/qwen3-tokenizer\
    prompt="Today is a beautiful day to"\
    model_name=<qwen3-30b-a3b|qwen3-235b-a22b|qwen3-480b-a35b>\
    per_device_batch_size=1\
    max_target_length=128\
    ici_fsdp_parallelism=256\
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/decode.sh (shellscript)
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --job-name=crowd-pilot_qwen3-0.6b_decode
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# Log the sbatch script
cat $0
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/decode.sh (shellscript)
<act focus file="sbatch_scripts/coinrun/train_tokenizer_coinrun.sh" />
// observation: file=sbatch_scripts/coinrun/train_tokenizer_coinrun.sh
```shellscript
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate



export SLURM_NTASKS=1

# ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared'
# job_name=$SLURM_JOB_NAME
# slurm_job_id=$SLURM_JOB_ID
# data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp'

# CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name_$slurm_job_id
# mkdir -p $CHECKPOINT_DIR

# # echo 'starting training'
# # srun echo 'hello'

# python train_tokenizer.py \
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/decode.sh (shellscript)
PROMPT=$(cat << 'EOF'
<act focus file="sbatch_scripts/coinrun/train_tokenizer_coinrun.sh" />
// observation: file=sbatch_scripts/coinrun/train_tokenizer_coinrun.sh
```shellscript
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate



export SLURM_NTASKS=1

# ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared'
# job_name=$SLURM_JOB_NAME
# slurm_job_id=$SLURM_JOB_ID
# data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp'

# CHECKPOINT_DIR=$ws_dir/checkpoints/$job_name_$slurm_job_id
# mkdir -p $CHECKPOINT_DIR

# # echo 'starting training'
# # srun echo 'hello'

# python train_tokenizer.py \
EOF
)
[tab] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/prompt.txt (plaintext)
<act focus file="sbatch_scripts/coinrun/train_tokenizer_coinrun.sh" />
observation: file=sbatch_scripts/coinrun/train_tokenizer_coinrun.sh
```shellscript
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate



export SLURM_NTASKS=1

ws_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared
job_name=SLURM_JOB_NAME
slurm_job_id=SLURM_JOB_ID
data_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp

CHECKPOINT_DIR=ws_dir/checkpoints/job_name_slurm_job_id
mkdir -p CHECKPOINT_DIR

echo starting training
srun echo hello

python train_tokenizer.py
```
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/decode.sh (shellscript)
<act focus file="sbatch_scripts/coinrun/train_tokenizer_coinrun.sh" />
observation: file=sbatch_scripts/coinrun/train_tokenizer_coinrun.sh
```shellscript
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate



export SLURM_NTASKS=1

ws_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared
job_name=SLURM_JOB_NAME
slurm_job_id=SLURM_JOB_ID
data_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_tmp

CHECKPOINT_DIR=ws_dir/checkpoints/job_name_slurm_job_id
mkdir -p CHECKPOINT_DIR

echo starting training
srun echo hello

python train_tokenizer.py
```
[tab] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/maxtext_decode.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --job-name=crowd-pilot_qwen3-0.6b_decode
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# Log the sbatch script
cat $0

cd maxtext
source .venv/bin/activate

CHECKPOINT_PATH=/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/outputs/crowd-pilot_qwen3-0.6b_batch_size_4_32858/checkpoints/0/items
PROMPT=$(cat << 'EOF'
<act focus file="sbatch_scripts/coinrun/train_tokenizer_coinrun.sh" />
// observation
EOF
)
MODEL_NAME=qwen3-0.6b
PER_DEVICE_BATCH_SIZE=1
MAX_TARGET_LENGTH=128

python3 -m MaxText.decode src/MaxText/configs/base.yml\
    load_parameters_path=$CHECKPOINT_PATH\
    tokenizer_path=src/MaxText/assets/qwen3-tokenizer\
    prompt="$PROMPT"\
    model_name=$MODEL_NAME\
    per_device_batch_size=$PER_DEVICE_BATCH_SIZE\
    max_target_length=$MAX_TARGET_LENGTH\
    skip_jax_distributed_system=True
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/vllm_decode.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=8
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --job-name=crowd-pilot_qwen3-0.6b_decode
#SBATCH --requeue
#SBATCH --signal=b:usr1@300 # 5 min before timeout

# Log the sbatch script
cat $0

cd maxtext
source .venv/bin/activate

CHECKPOINT_PATH=/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/outputs/crowd-pilot_qwen3-0.6b_batch_size_4_32858/checkpoints/0/items
PROMPT=$(cat << 'EOF'
<act focus file="sbatch_scripts/coinrun/train_tokenizer_coinrun.sh" />
// observation
EOF
)
MODEL_NAME=qwen3-0.6b
PER_DEVICE_BATCH_SIZE=1
MAX_TARGET_LENGTH=128

python3 -m MaxText.decode src/MaxText/configs/base.yml\
    load_parameters_path=$CHECKPOINT_PATH\
    tokenizer_path=src/MaxText/assets/qwen3-tokenizer\
    prompt="$PROMPT"\
    model_name=$MODEL_NAME\
    per_device_batch_size=$PER_DEVICE_BATCH_SIZE\
    max_target_length=$MAX_TARGET_LENGTH\
    skip_jax_distributed_system=True
[tab] /fast/home/franz.srambical/crowd-pilot/maxtext/src/MaxText/integration/tunix/tunix_adapter.py (python)
# Copyright 2023–2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Adapter for integrating MaxText Transformer models with Tunix.

This module provides the `TunixMaxTextAdapter` class, which wraps a MaxText
Transformer model to expose a call signature compatible with Tunix Trainers.
It also handles weight mapping for compatibility with Hugging Face models.
"""

from __future__ import annotations

from typing import Optional, Tuple, Any

from jax import Array
from flax import nnx
from MaxText.layers.models import Transformer
from MaxText.integration.tunix.weight_mapping import VLLM_WEIGHT_MAPPING


class TunixMaxTextAdapter(nnx.Module):
  """Adapter exposing Tunix Trainer call signature over a Transformer model."""

  def __init__(
      self,
      base_model: Transformer,
  ):
    super().__init__()
    self.base = base_model

  # ------------------------------------------------------------------ #
  # Tunix call signature
  # ------------------------------------------------------------------ #
  def __call__(
      self,
      input_tokens: Array,  # [B, L]
      positions: Array,  # [B, L]
      cache: Optional[Any],  # Tunix currently passes None from Trainers
      attention_mask: Optional[Array],  # [B, L, L] or None
      output_hidden_states: bool = False,  # ignored
  ) -> Tuple[Array, None]:
    """Forward compatible with Tunix Trainers default loss.
    Returns logits, None.
    """
    logits = self.base(
        decoder_input_tokens=input_tokens,
        decoder_positions=positions,
        # TODO: @mazumdera - add support for packing
        decoder_segment_ids=None,
    )
    return logits, None

  def to_hf_mappings(self):
    return VLLM_WEIGHT_MAPPING[self.base.config.model_name].to_hf_mapping()

  def to_hf_transpose_keys(self):
    return VLLM_WEIGHT_MAPPING[self.base.config.model_name].to_hf_transpose_keys()

  def to_hf_hook_fns(self):
    return VLLM_WEIGHT_MAPPING[self.base.config.model_name].to_hf_hook_fns()

  def lora_to_hf_mappings(self):
    return VLLM_WEIGHT_MAPPING[self.base.config.model_name].lora_to_hf_mappings()
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/convert_checkpoint_to_hf.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --time=02:00:00
#SBATCH --cpus-per-task=4
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/%x_%j.log
#SBATCH --job-name=convert_maxtext_to_hf
#SBATCH --requeue

# Usage:
#   sbatch convert_checkpoint_to_hf.sh \
#     /path/to/maxtext/checkpoints/0/items \
#     qwen3-0.6b \
#     /path/to/output/hf_dir \
#     <optional_hf_token>

set -euo pipefail

cat $0

if [ $# -lt 3 ]; then
  echo "Usage: $0 <CHECKPOINT_PATH> <MODEL_NAME> <HF_OUT_DIR> [HF_ACCESS_TOKEN]" >&2
  exit 1
fi

CHECKPOINT_PATH="$1"
MODEL_NAME="$2"
HF_OUT_DIR="$3"
HF_ACCESS_TOKEN="${4-}"

cd /home/franz.srambical/crowd-pilot/maxtext
source .venv/bin/activate

mkdir -p "$HF_OUT_DIR"

python3 -m MaxText.utils.ckpt_conversion.to_huggingface src/MaxText/configs/base.yml \
  model_name=$MODEL_NAME \
  load_parameters_path=$CHECKPOINT_PATH \
  base_output_directory=$HF_OUT_DIR \
  scan_layers=false \
  use_multimodal=false \
  hf_access_token=$HF_ACCESS_TOKEN

echo "Converted HF model saved to: $HF_OUT_DIR"
[content] slurm/dev/franz/berlin/crowd-pilot/qwen_0_6/vllm_decode.sh (shellscript)
PROMPT_FILE="$SCRIPT_DIR/prompt.txt"

# Run the Python script
python3 "$PYTHON_SCRIPT" \
    --model_path "$HF_CHECKPOINT_PATH" \
    --prompt_file "$PROMPT_FILE" \
    --max_tokens 128 \
    --temperature 0.0 \
    --top_p 1.0 \
    --top_k 0 \
    --dtype float16