Edit Flows base model

Training

Trained on a single 8xH100 node:

accelerate launch \
 --config_file scripts/accelerate_configs/zero2.yaml \
 examples/editflow/llada/adapt.py --model_name_or_path "GSAI-ML/LLaDA-8B-Instruct" \
 --lm_head_key "model.transformer.ff_out" \
 --init_editflow_from_src True \
 --per_device_train_batch_size 1 \
 --per_device_eval_batch_size 1 \
 --gradient_accumulation_steps 4 \
 --dataset_args "allenai/tulu-3-sft-mixture[train:500000]|lamm-mit/bio-silk-mech-mix-q-a-35K-messages-only|lamm-mit/graph_reasoning_v3_messages" \
 --output_dir "models/LlaDA-8B-EditFlow-instruct-v500" \
 --x0_sampler "masks[length:128]" --max_length 1500 \
 --num_train_epochs 4 \
 --learning_rate 1e-5 \
 --push_to_hub True --save_strategy "steps" --save_steps 1000 \
 --hub_model_id lamm-mit/LlaDA-8B-EditFlow-instruct-v500 \
 --hub_private_repo True --eval_strategy "no" \
 --warmup_steps 50 

Sampling

python examples/editflow/sample.py \
 --model_name_or_path "odels/LlaDA-8B-EditFlow-instruct-v500" \
 --mask_length 128 --seed 7070 \
 --prompt "Define materiomics."

Base model