8000 [Clean] Fix warnings and abnormal outputs during training by lmyybh · Pull Request #393 · sgl-project/SpecForge · GitHub
[go: up one dir, main page]
Skip to content

[Clean] Fix warnings and abnormal outputs during training#393

Open
lmyybh wants to merge 4 commits intosgl-project:mainfrom
lmyybh:clean
Open

[Clean] Fix warnings and abnormal outputs during training#393
lmyybh wants to merge 4 commits intosgl-project:mainfrom
lmyybh:clean

Conversation

@lmyybh
Copy link
@lmyybh lmyybh commented Dec 26, 2025 •
edited
Loading

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

@lmyybh
[Clean] Fix warnings and abnormal outputs
054c8cf 
- Update FSDP checkpoint API to deprecated warnings
- Fix distributed init order to avoid device binding warnings
- Fix tqdm progress bar display on non-rank-0 processes
- Replace logger with print to avoid logging warnings
@lmyybh
[Improve] Add eval logging and fix model state after evaluation
f9ee3c9
@gemini-code-assist
Copy link
Contributor
gemini-code-assist bot commented Dec 26, 2025

Warning

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

@FrankLeeeee
Copy link
Collaborator
FrankLeeeee commented Dec 26, 2025

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

@lmyybh
Copy link
Author
lmyybh commented Dec 26, 2025

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

Done

@sleepcoo
Copy link
Collaborator
sleepcoo commented Dec 29, 2025

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 lmyybh requested a review from FrankLeeeee as a code owner December 29, 2025 09:49
@lmyybh
Copy link
Author
lmyybh commented Dec 29, 2025

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.

uygnef
uygnef reviewed Feb 4, 2026
View reviewed changes
specforge/utils.py

def print_with_rank(message):
if dist.is_available() and dist.is_initialized():
logger.info(f"rank {dist.get_rank()}: {message}")
Copy link
Collaborator
@uygnef uygnef Feb 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think using logger.info is preferable. To enable it, simply add logging.getLogger().setLevel(logging.INFO) in your main function.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@uygnef uygnef uygnef left review comments

@shuaills shuaills Awaiting requested review from shuaills shuaills is a code owner

@sleepcoo sleepcoo Awaiting requested review from sleepcoo sleepcoo is a code owner

@FlamingoPg FlamingoPg Awaiting requested review from FlamingoPg FlamingoPg is a code owner

@FrankLeeeee FrankLeeeee Awaiting requested review from FrankLeeeee FrankLeeeee is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@lmyybh @FrankLeeeee @sleepcoo @uygnef
0