[Clean] Fix warnings and abnormal outputs during training

[lmyybh]

Motivation

During training, I noticed several warnings and abnormal outputs cluttering the logs:

1. FSDP deprecation warnings when saving checkpoints using the old state_dict_type() context manager
2. NCCL initialization warnings due to incorrect device binding order in distributed setup
3. print_on_rank0 not producing any output due to improper logger configuration
4. Missing model state reset after evaluation, which could lead to unexpected behavior in subsequent training steps

Modifications

1. Checkpoint Saving: Replace deprecated state_dict_type with get_model_state_dict
2. Distributed Initialization: Pass device_id parameter to suppress NCCL warnings
3. print_on_rank0: Replace logger.info() with print()
4. Evaluation:
   • Call draft_model.train() after evaluation to ensure correct model state
   • Show progress bar only on rank 0
   • Print metrics in record_metrics only during evaluation
5. Dependencies: Add tensorboard package
6. Replace Optimizer with BF16Optimizer

Related Issues

Accuracy Test

Benchmark & Profiling

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[FrankLeeeee]

Hi, can you apply pre-commit for your code?

[lmyybh]

Hi, can you apply pre-commit for your code?

Done

[sleepcoo]

can u fix unit test？

  File "/__w/SpecForge/SpecForge/specforge/distributed.py", line 76, in init_distributed
    local_rank = int(os.environ["LOCAL_RANK"])
                     ~~~~~~~~~~^^^^^^^^^^^^^^
  File "<frozen os>", line 679, in __getitem__
KeyError: 'LOCAL_RANK'

[lmyybh]

can u fix unit test？

  File "/__w/SpecForge/SpecForge/specforge/distributed.py", line 76, in init_distributed
    local_rank = int(os.environ["LOCAL_RANK"])
                     ~~~~~~~~~~^^^^^^^^^^^^^^
  File "<frozen os>", line 679, in __getitem__
KeyError: 'LOCAL_RANK'

Done. And I will continuously check the unit test results and fix the errors.