Bucket-Level MOO: Multilingual Fine-Tuning via Localized Gradient Conflict Resolution

Official implementation of "Multilingual Fine-Tuning via Localized Gradient Conflict Resolution".

Repository structure

bucket_moo.py          # Core BK MOO components (ZeRO-2 patch, callbacks)
bucket_moo_train.py    # Entry point: BK MOO training (ZeRO-2)
global_moo.py          # Core Global MOO components (Gram-matrix solvers, DDP trainer)
global_moo_train.py    # Entry point: Global MOO baseline (standard DDP)
vsft_train.py          # Entry point: Vanilla SFT baseline
min_norm_solvers.py    # MGDA min-norm solver
dataloader.py          # Task-parallel multilingual datasets / collators
utils.py               # Languages, prompts, tokenizer setup, data processing
eval/
  eval_utility.py      # Multilingual evaluation harness (SGLang offline engine)
  grader.py, parser.py # Math/MCQ answer extraction and grading
acc_config/            # Accelerate launch configs (ddp_8, ds_zero2_8)
exp_config/            # Training hyperparameters (st_config_llama3, st_config_qwen3)

MOO modes

Both bucket_moo_train.py and global_moo_train.py expose the same --mode flag:

--mode	Method
pc	PCGrad — project out pairwise conflicting components, then sum
mgda	MGDA — min-norm element of the convex hull of task gradients
mgda_l2	MGDA on L2-normalized task gradients (equal-magnitude variant)
cagrad	CAGrad — conflict-averse direction near the average gradient

Installation

Use one of the pinned environment files:

conda create -n bkmoo python=3.12 -y
conda activate bkmoo
pip install -r requirements.txt

Evaluation additionally requires SGLang for the offline inference engine (install separately to match your CUDA setup).

We used 8 H200 GPUs by default (one language per GPU). The acc_config/*.yaml files assume num_processes: 8; adjust them for a different device count.

Data

Training uses the lima_s1 mixture (LIMA + s1 instructions, per language) from Hugging Face iNLP-Lab/multilingual-lima and iNLP-Lab/multilingual-s1 (prompt / output columns). Eight in-distribution languages are used for fine-tuning:

ar, bn, en, id, it, ko, sw, zh

Evaluation additionally reports out-of-distribution languages (de, es, fr, ja, pt).

Training

Bucket-Level MOO (DeepSpeed ZeRO-2):

accelerate launch --config_file acc_config/ds_zero2_8.yaml bucket_moo_train.py \
    --mode cagrad --base meta-llama/Meta-Llama-3-8B \
    --alpha 0.5 --rescale 1 \
    --training_config exp_config/st_config_llama3.yaml

Global MOO baseline (standard DDP, full-gradient MOO):

accelerate launch --config_file acc_config/ddp_8.yaml global_moo_train.py \
    --mode cagrad --base Qwen/Qwen3-4B-Base \
    --training_config exp_config/st_config_qwen3.yaml

Vanilla SFT baseline:

accelerate launch --config_file acc_config/ds_zero2_8.yaml vsft_train.py \
    --base meta-llama/Meta-Llama-3-8B \
    --training_config exp_config/st_config_llama3.yaml

CAGrad-only flags: --alpha (trust-region radius) and --rescale (0=none, 1=/(1+alpha^2), 2=/(1+alpha)).

Evaluation

eval/eval_utility.py runs an SGLang offline engine over four multilingual benchmarks (PolyMath math_easy, Global-MMLU-Lite, Belebele, MuBench ARC-Easy), then grades the predictions and writes a per-task report plus summary.json.

python -m eval.eval_utility \
    --model_name saved_models/cagrad/lima_s1/meta-llama/Meta-Llama-3-8B \
    --dp_size 8                 # add --zero_shot for zero-shot CoT prompting

Citation

@misc{hoang2026multilingualfinetuninglocalizedgradient,
      title={Multilingual Fine-Tuning via Localized Gradient Conflict Resolution}, 
      author={Long P. Hoang and Yiran Zhao and Wei Lu and Wenxuan Zhang},
      year={2026},
      eprint={2606.05613},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2606.05613}, 
}