p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/dev/franz/berlin/crowd-pilot/train_maxtext_array_record_10_epochs.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_maxtext_array_record_10_epochs
#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/
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

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=50000 \
  num_epoch=10 \
  skip_jax_distributed_system=True
[tab] slurm/dev/franz/berlin/crowd-pilot/train_maxtext_array_record.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_maxtext_array_record_dev
#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/
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

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=50000 \
  skip_jax_distributed_system=True
[tab] slurm/dev/franz/berlin/crowd-pilot/train_maxtext_parquet_10_epochs.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_maxtext_10_epochs
#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/
GRAIN_TRAIN_FILES="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/parquet/train_*.parquet"
GRAIN_EVAL_FILES="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/parquet/val_*.parquet"

env | grep SLURM

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=parquet \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=50000 \
  num_epoch=10 \
  skip_jax_distributed_system=True
[tab] slurm/dev/franz/berlin/crowd-pilot/train_maxtext_parquet.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_maxtext_dev
#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/
GRAIN_TRAIN_FILES="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/parquet/train_*.parquet"
GRAIN_EVAL_FILES="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/parquet/val_*.parquet"

env | grep SLURM

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=parquet \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=50000 \
  skip_jax_distributed_system=True
[tab] slurm/dev/franz/berlin/crowd-pilot/generate_array_record_dataset.sh (shellscript)
#!/bin/bash

set -uex

OUTPUT_DIR="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/array_record/"
CSV_ROOT="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/crowd-code-0.1/csv/"

uv run crowd-pilot/serialize_dataset_array_record.py --csv_root=$CSV_ROOT --output_dir=$OUTPUT_DIR
[tab] slurm/dev/franz/berlin/crowd-pilot/random-weights/train_maxtext_array_record_10_epochs.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_maxtext_array_record_10_epochs
#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/
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

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=25000 \
  num_epoch=10 \
  skip_jax_distributed_system=True
[tab] slurm/dev/franz/berlin/crowd-pilot/random-weights/train_maxtext_array_record.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_maxtext_array_record_dev
#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/
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

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=25000 \
  skip_jax_distributed_system=True
[content] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/train_maxtext_array_record.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_maxtext_array_record_dev
#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/
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

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=25000 \
  skip_jax_distributed_system=True
[content] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/train_maxtext_array_record_10_epochs.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_maxtext_array_record_10_epochs
#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/
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

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=25000 \
  num_epoch=10 \
  skip_jax_distributed_system=True
[content] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/convert_checkpoint.sh (shellscript)
set -uex

source .venv/bin/activate

MODEL_NAME="Qwen/Qwen3-Coder-30B-A3B-Instruct"
OUTPUT_DIR="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/checkpoint_pretrained/${MODEL_NAME}"

python3 -m MaxText.convert_qwen3_moe \
  --base_model_path ${OUTPUT_DIR} \
  --maxtext_model_path ${OUTPUT_DIR} \
  --model_size qwen3-coder-30b-a3b-instruct
[content] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/convert_checkpoint.sh (shellscript)
python3 -m MaxText.convert_qwen3_moe\
  --base_model_path /path/to/qwen3_hf_checkpoint\
  --maxtext_model_path gs://your-gcs-bucket/qwen3_maxtext_ckpt\
  --model_size <qwen3-30b-a3b|qwen3-235b-a22b|qwen3-480b-a35b>
[content] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/convert_checkpoint.sh (shellscript)
#!/usr/bin/env bash
set -uex

source .venv/bin/activate

MODEL_NAME="Qwen/Qwen3-Coder-30B-A3B-Instruct"
OUTPUT_DIR="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/checkpoint_pretrained/${MODEL_NAME}"

python3 -m MaxText.convert_qwen3_moe \
  --base_model_path ${OUTPUT_DIR} \
  --maxtext_model_path ${OUTPUT_DIR} \
  --model_size qwen3-coder-30b-a3b-instruct
[tab] maxtext/base_requirements/requirements.txt (pip-requirements)
absl-py
aqtp
array-record
cloud-accelerator-diagnostics
cloud-tpu-diagnostics
datasets
flax
gcsfs
google-api-python-client
google-cloud-aiplatform
google-cloud-monitoring
grain[parquet]
huggingface_hub
jax!=0.7.1, !=0.7.2
jaxlib!=0.7.1, !=0.7.2
jaxtyping
jsonlines
ml-collections
ml-goodput-measurement
numpy
omegaconf
optax
orbax-checkpoint
pathwaysutils
pillow
pre-commit
protobuf
pyink
pylint
pytest
pytype
sentencepiece
tensorboard-plugin-profile
tensorboardx
tensorflow-datasets
tensorflow-text
tensorflow
tiktoken
transformers
qwix
google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
[tab] maxtext/pyproject.toml (plaintext)
[build-system]
requires = ["hatchling", "hatch-requirements-txt"]
build-backend = "hatchling.build"

[tool.hatch.version]
path = "src/MaxText/__init__.py"

[project]
name = "maxtext"
description = "MaxText is a simple, performant and scalable Jax LLM!"
dynamic = ["version", "optional-dependencies"]
requires-python = ">=3.12"
readme = "README.md"
license = "Apache-2.0"
keywords = ["llm", "jax", "llama", "mistral", "mixtral", "gemma", "deepseek"]
classifiers = [
    "Development Status :: 4 - Beta",
    "Intended Audience :: Developers",
    "Intended Audience :: Science/Research",
    "Programming Language :: Python",
    "Programming Language :: Python :: 3.12",
    "Programming Language :: ML",
    "Environment :: GPU :: NVIDIA CUDA :: 12",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = []

[tool.hatch.metadata.hooks.requirements_txt.optional-dependencies]
tpu = ["generated_requirements/tpu-requirements.txt"]
cuda12 = ["generated_requirements/cuda12-requirements.txt"]

[project.urls]
Repository = "https://github.com/AI-Hypercomputer/maxtext.git"
"Bug Tracker" = "https://github.com/AI-Hypercomputer/maxtext/issues"

[tool.hatch.metadata]
allow-direct-references = true

[tool.hatch.build.targets.wheel]
packages = ["src/MaxText", "src/install_maxtext_extra_deps"]

[project.scripts]
install_maxtext_github_deps = "install_maxtext_extra_deps.install_github_deps:main"
[content] maxtext/src/MaxText/convert_qwen3_moe.py (python)

        "num_hidden_layers": 40,
        "num_attention_heads": 40,
        "num_key_value_heads": 4,
        "hidden_size": 4096,
        "head_dim": 128,
        "num_experts": 128,
        "moe_intermediate_size": 1536,
    },
[tab] maxtext/src/MaxText/configs/models/qwen3-30b-a3b.yml (yaml)
# Copyright 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.

# Model config for Qwen3-30B-A3B

# Core Architectural Parameters
decoder_block: "qwen3_moe"
base_emb_dim: 2048
base_mlp_dim: 768
base_num_query_heads: 32
base_num_kv_heads: 4
base_num_decoder_layers: 48
head_dim: 128
mlp_activations: ["silu", "linear"]
vocab_size: 151936
normalization_layer_epsilon: 1.0e-6
use_qk_norm: True

# MoE Specific Parameters
num_experts: 128
num_experts_per_tok: 8
base_moe_mlp_dim: 768
norm_topk_prob: true

# RoPE Settings
rope_max_timescale: 10_000_000

# General Model Settings
enable_dropout: False
[tab] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/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-30b-a3b_dev
#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-Coder-30B-A3B-Instruct/
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-30b-a3b \
  max_target_length=8192 \
  eval_interval=500 \
  steps=25000 \
  batch_size=1 \
  tokenizer_type=huggingface \
  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
[tab] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/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-30b-a3b_dev
#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-Coder-30B-A3B-Instruct/
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-30b-a3b \
  max_target_length=8192 \
  eval_interval=500 \
  steps=25000 \
  batch_size=1 \
  tokenizer_type=huggingface \
  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
[tab] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/train_10_epochs.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_maxtext_array_record_10_epochs
#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/
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

python3 -m MaxText.train \
  MaxText/configs/base.yml \
  run_name=${job_name}_${slurm_job_id} \
  base_output_directory=$BASE_OUTPUT_DIRECTORY \
  dataset_type=grain \
  grain_file_type=arrayrecord \
  grain_train_files=$GRAIN_TRAIN_FILES \
  grain_eval_files=$GRAIN_EVAL_FILES \
  eval_interval=500 \
  steps=25000 \
  num_epoch=10 \
  skip_jax_distributed_system=True
[content] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/train_10_epochs.sh (shellscript)
# 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-30b-a3b \
  max_target_length=8192 \
  eval_interval=500 \
  steps=25000 \
  batch_size=1 \
  tokenizer_type=huggingface \
  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
[tab] slurm/dev/franz/berlin/crowd-pilot/qwen-coder-small/convert_checkpoint_slurm.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --time=4:00:00
#SBATCH --cpus-per-task=16
#SBATCH --mem=300G
#SBATCH --output=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/convert_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/jafar_ws/logs/franz/crowd-pilot/maxtext/convert_%j.log
#SBATCH --job-name=convert_qwen3-30b-a3b
#SBATCH --requeue

# Log the sbatch script
cat $0

# Set higher host memory limit for large model conversion
export XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB=250
export TF_GPU_ALLOCATOR=cuda_malloc_async

cd maxtext
source .venv/bin/activate

MODEL_NAME="Qwen/Qwen3-Coder-30B-A3B-Instruct"
HF_MODEL_PATH="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/checkpoint_pretrained_hf/${MODEL_NAME}"
OUTPUT_DIR="/fast/project/HFMI_SynergyUnit/jafar_ws/data/crowd-pilot/checkpoint_pretrained_maxtext/${MODEL_NAME}"

python3 -m MaxText.convert_qwen3_moe\
  --base_model_path ${HF_MODEL_PATH}\
  --maxtext_model_path ${OUTPUT_DIR}\
  --model_size qwen3-30b-a3b