Skip to content

feat: enable Hugging Face eval_loss on sharded datasets - #777

Open
DanielAsadi wants to merge 1 commit into
NVIDIA:mainfrom
DanielAsadi:feat/hf-eval-loss
Open

DanielAsadi wants to merge 1 commit into
NVIDIA:mainfrom
DanielAsadi:feat/hf-eval-loss

Conversation

@DanielAsadi

Copy link
Copy Markdown

Summary

Hugging Face eval_strategy is documented on GR00T, but the sharded dataset path cannot actually produce eval_loss:

  • DatasetFactory.build() asserts eval_strategy == "no" and always returns eval_dataset=None.
  • The mixture iterator is infinite, so an eval pass would never finish.
  • HF Accelerate wraps the eval iterable dataloader and broadcasts CPU objects through NCCL.
  • Trainer.prediction_step drops the loss because Gr00tN1d7 has empty label_names and no return_loss argument, so eval emits eval_runtime but no eval_loss.

This change keeps training behavior when eval_strategy is "no". When it is not:

  • split episodes deterministically into disjoint train/eval sets
  • iterate eval shards once, padding the schedule to world_size * num_workers
  • return a raw eval DataLoader (same as train) and put the processor in eval mode
  • always return the model forward loss from prediction_step

Test plan

  • tests/gr00t/data/test_dataset_factory.py
  • tests/gr00t/data/test_sharded_datasets.py
  • tests/gr00t/experiment/test_trainer_eval_dataloader.py
  • tests/gr00t/experiment/test_trainer_prediction_step_loss.py
  • 1-GPU fine-tune smoke: eval ran at steps 10 and 20, eval_loss appeared in trainer state / logs, training resumed, checkpoint written
  • 4-GPU smoke still pending (exercises shard padding / multi-rank collectives)

Gr00tN1d7 already returns a training loss, but DatasetFactory refused to
build an eval set, the mixture iterator never terminated, Accelerate
broadcast CPU objects through NCCL, and Trainer.prediction_step dropped
the loss because the model has empty label_names. Split episodes, make
eval iteration finite and rank-aligned, skip Accelerate dispatch, and
return the forward loss so Hugging Face logs eval_loss.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant