NLLB-200 Chinese-to-English Translation Fine-tuning

LoRA fine-tuning of facebook/nllb-200-distilled-1.3B for Chinese-to-English translation (zho_Hans -> eng_Latn) on WMT19 zh-en.

Chinese docs: README.zh-CN.md

Course assignment: HKU DASC7606A-B Assignment 2
Original assignment repo: https://github.com/hkukend/DASC7606A-B-A2

---

Result Summary

Metric / Constraint	Result
Test BLEU	24.9
Metric	sacreBLEU
Test split	WMT19 zh-en validation
Test rows	3,981
Training data	600,000 sentence pairs
Validation data	2,000 sentence pairs
Hardware	1 x RTX 4080S 16 GB
Training time	about 5.5 hours
Time limit	under 12 hours

The reported score is measured on the original WMT19 validation split used as the test set. The test set is not filtered, shuffled, or modified; row count, byte size, and dataset fingerprint are checked before training.

---

Technical Highlights

Area	Implementation
Base model	NLLB-200 distilled 1.3B
Parameter-efficient tuning	LoRA adapters instead of full fine-tuning
Adapter coverage	Attention projections and feed-forward projections
Data quality	Multi-stage training-only cleaning and exact-pair deduplication
Decoding	Forced English BOS token + 5-beam search
Model selection	Best checkpoint selected by validation BLEU
Evaluation	Fixed sacreBLEU metric implementation
Reproducibility	Fixed seed, deterministic CUBLAS workspace, test-set integrity check

---

Pipeline

flowchart LR
    A[WMT19 zh-en] --> B[Shuffle with seed 42]
    B --> C[600k train pairs]
    B --> D[2k validation pairs]
    A --> E[Original validation split as test]
    C --> F[Normalize, filter, deduplicate]
    D --> G[Basic structure check]
    E --> H[Test integrity check]
    F --> I[NLLB tokenizer]
    G --> I
    H --> I
    I --> J[NLLB-200 distilled 1.3B + LoRA]
    J --> K[Seq2SeqTrainer]
    K --> L[Best checkpoint by validation BLEU]
    L --> M[Test sacreBLEU: 24.9]

Loading

---

Deep Dive

1. Model Strategy

The project uses facebook/nllb-200-distilled-1.3B as a strong multilingual base model and trains only LoRA adapters. This keeps the number of trainable parameters small while still adapting the model to the Chinese-to-English domain.

LoRA setting	Value
Rank r	32
lora_alpha	64
lora_dropout	0.1
Target modules	q_proj, k_proj, v_proj, o_proj, fc1, fc2
Bias	none
Task type	SEQ_2_SEQ_LM

The tokenizer is initialized with source language zho_Hans and target language eng_Latn. During generation, forced_bos_token_id is set to the English language token so NLLB always decodes into English.

2. Data Strategy

The raw corpus is WMT19 zh-en. The training split is shuffled with seed=42, then sliced into 600,000 training pairs and 2,000 validation pairs. The WMT19 validation split is reserved as the final test set.

Training data receives the full cleaning pipeline in dataset.py:

• Text normalization: remove Unicode control characters, convert fullwidth ASCII to halfwidth, apply NFC normalization, and collapse whitespace.
• Length and ratio filtering: keep sentence pairs with 4-256 characters per side and a zh/en length ratio within [0.2, 5.0].
• Language consistency filtering: require Chinese text on the source side and reject target sentences with too much Chinese text.
• Noise filtering: remove copy-like pairs, repeated spans, excessive special characters, numeric-heavy samples, and whitespace-heavy samples.
• Deduplication: remove exact duplicate translation pairs.

Validation data only receives a lightweight structure check. The test set is kept unchanged and verified by not_change_test_dataset() in utils.py.

3. Training Strategy

Training uses Hugging Face Seq2SeqTrainer for 1 epoch. The configuration is tuned for a single 16 GB GPU and completed on one RTX 4080S 16 GB in about 5.5 hours, comfortably below the 12-hour assignment limit.

Hyperparameter	Value
Epochs	1
Per-device train batch size	24
Gradient accumulation	3
Effective batch size	about 72
Learning rate	1e-5
Scheduler	cosine
Warmup ratio	0.1
Weight decay	0.01
Label smoothing	0.1
Max grad norm	1.0
Eval / save interval	800 steps
Precision	bf16 when supported, otherwise fp16

The best checkpoint is selected with load_best_model_at_end=True using validation BLEU as the selection metric.

4. Decoding and Evaluation

Generation uses 5-beam search with generation_max_length=160 and length_penalty=1.2. Evaluation uses the fixed compute_metrics function in evaluation.py, which reports sacreBLEU.

---

Project Structure

constants.py            # Model checkpoint, output directory, max input/target length
model.py                # NLLB tokenizer/model loading and LoRA adapter setup
dataset.py              # WMT19 loading, cleaning, validation checks, tokenization
trainer.py              # Seq2SeqTrainingArguments and SafeSeq2SeqTrainer
evaluation.py           # sacreBLEU metric computation
utils.py                # Random seeds, CUDA env setup, test-set integrity check
main.py                 # Train and evaluate end-to-end
evaluate_checkpoint.py  # Evaluate saved LoRA checkpoints on the test set
pyproject.toml          # uv-managed dependencies

---

Setup

Requirements:

• Python >= 3.13
• CUDA GPU, preferably with bf16 support
• uv

Install dependencies:

uv sync

---

Usage

Train and evaluate:

uv run python main.py

Evaluate the latest saved checkpoint:

uv run python evaluate_checkpoint.py

Evaluate a specific checkpoint:

uv run python evaluate_checkpoint.py results/checkpoint-10000

Evaluate all matching checkpoints:

uv run python evaluate_checkpoint.py 'results/checkpoint-*'

---

Output

Training writes checkpoints to results/checkpoint-*. After training, main.py evaluates the selected model on the unchanged test set and prints metrics in this form:

Test Metrics: {'test_bleu': 24.9, ...}

---

Reproducibility

The project fixes random, numpy, and torch seeds to 42, configures CUBLAS workspace determinism, and checks the WMT19 test split before evaluation. cuDNN benchmarking remains enabled for speed, so exact bitwise reproducibility is not guaranteed across different machines or driver/library versions.