Feat: SGL backend for online SD training

[h-guo18]

What does this PR do?

Type of change: new feature

Overview:

• New trainer with different base model backends available for online training.

  • Now the base model can be SGlang, HF-TP, while the student model DDP on different devices.
  • Trainer based on previous experimental PR352.
  • SGL wrapper based on specforge PR239;

• Other improvements:

  • Move train_acc.item() out from eagle forward to avoid cuda graph break during torch compile;
  • Optionally turn off ar validation pbar to avoid log overflow.

Usage

export NCCL_CUMEM_ENABLE=1;
python train.py  \
 --out_path $OUT \
 --data_path $DATA \
 --model_path $MODEL \
 --teacher_backend <"sglang" or "hf"> \
 --teacher_devices 0,1,2,3 \
 --student_devices 4,5,6,7 \
 --teacher_ep_size 1  #sglang backend only

Testing

Parallelism

• Tested base model TP on HF backend;
• Tested base model TP and EP on SGLang backend;

Training Quality Test

Compared previous HF trainer, new trainer-HF backend and trainer-SGL backend;
Setting: Llama3.2-1B, magpie, bs=8, lr1e-4, seqlen1k;

• slight different is a result of:
  • different data loading random seed
  • sgl backend use bf16 lm head, instead of fp32;

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.45%. Comparing base (1aaa77d) to head (70515dc).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #564   +/-   ##
=======================================
  Coverage   74.45%   74.45%           
=======================================
  Files         182      182           
  Lines       18250    18250           
=======================================
  Hits        13588    13588           
  Misses       4662     4662           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[yeyu-nvidia]

Is this code borrowed from OSS? if so, you will need to follow the OSS code procedure to file a ticket and get approval

[yeyu-nvidia]

According to the figures, there is a big accuracy/loss gap between SGL trainer and our current trainer. We need to figure it out before this can be merged. Also, as this PR may change our training accuracy fundamentally, we need to train larger models than llama 1B