Bugfixes - Custom hooks; optimizer warmup as epochs; dataset cache as a separate configuration; validation at end of epoch fix #505
Conversation
This commit addresses feedback from the cosmos-rl fork by enhancing the
hook system and providing examples for custom logging integration:
1. Added training lifecycle hooks to SFTPolicyWorker:
- pre_training_hook: Called before training loop starts
- pre_training_step_hook: Called before each training step
- post_training_step_hook: Called after each training step
- post_training_hook: Called after training completes
2. Improved setup_hooks() documentation with detailed descriptions of
all available hooks (training and validation)
3. Added custom_loggers_and_hooks.py with:
- create_status_logger(): Factory for custom logging functions
- create_validation_hooks(): Factory for validation lifecycle hooks
- create_training_hooks(): Factory for training lifecycle hooks
- create_all_hooks(): Convenience function for all hooks
- TAOStatusLogger: Full TAO-compatible logger that writes to status.json
(uses nvidia_tao_core.loggers if available, falls back to JSON)
- _FallbackStatusLogger: JSON file logger when TAO core unavailable
4. Updated README.md with comprehensive documentation on:
- Custom logger function signature and usage
- Available hooks and their report_data contents
- Example integrations with external monitoring systems
These changes enable users to integrate external logging systems
(e.g., TAO, MLflow) via custom_logger_fns and hook_fns parameters
passed through cosmos_rl.launcher.worker_entry, eliminating the need
for hardcoded logging logic in the trainer.
The TAOStatusLogger writes to:
{TAO_API_RESULTS_DIR}/{TAO_API_JOB_ID}/status.json
in the same format as the original log_tao_status function.
This script demonstrates how to use custom_logger_fns and hook_fns
for TAO-compatible status logging in SFT training.
Key features:
- TAOStatusLogger integration that writes to status.json
- Custom dataset support for llava-format data
- Separate training and validation dataset support
- Configurable via [custom] section in TOML config
Usage:
cosmos-rl --config spec.toml scripts/custom_sft.py
With TAO logging enabled (set TAO_API_JOB_ID env var):
TAO_API_JOB_ID=job123 cosmos-rl --config spec.toml scripts/custom_sft.py
The script passes custom_logger_fns and hook_fns to worker_entry.main()
instead of using the @monitor_status decorator, providing a cleaner
separation of concerns.
There was a problem hiding this comment.
Overall LGTM. Since the previous custom_example folder is reserved for the purpose of applying custom policy/rollout backend. We might move the newly added files in this folder into another new folder such as tools/custom_hooks?
Also, pre-commit run --all-files might be needed before commit to pass the lint check?
ramanathan831
force-pushed
the
feature/enhanced-hooks-and-custom-loggers
branch
from
83bc016 to
91e72b2
Compare
ramanathan831
force-pushed
the
feature/enhanced-hooks-and-custom-loggers
branch
from
91e72b2 to
b943dc0
Compare
If vllm backend is used in rollout but no vllm installed, an unsupported message will appear.
1. Optimize the enqueue client timeout 2. Change the multi-modal data type from fp32 to fp16 3. Fix the lepton endpoint multi-replicas round-robin issue 4. Add latency benchmark 5. Support separated scores redis service (optional, default is still colocated)
Co-authored-by: liang.feng <liangf@nvidia.com>
…os#508) Support colocated mode of Policy and Rollout on the same devices but in separated processes. It must: 1. share all the same devices. That means: one policy replica and one rollout replica must have the same world size. 2. the model is not huge that policy and rollout could exist on the same devices. 3. `gpu_memory_utilization` could be lower to decrease the memory vLLM occupies.
ramanathan831
force-pushed
the
feature/enhanced-hooks-and-custom-loggers
branch
from
aaa184f to
d84faa5
Compare
Great job for these tests! Yeah, regarding this part, I have left a new comment about how to calculate the number od total samples and the whole val_avg_loss. We could check whether the point in that comment makes sense? |
Thx for the explanation! Yeah, then I think that would be okay to just use the |
- sft_worker.py: Replace all_reduce sample aggregation with len(val_data_loader.dataset) which already gives total samples across all ranks. Remove unused torch.distributed import. - sft_trainer.py: Change per-rank validation log from info to debug level to reduce log noise in production.
Resolve conflicts in sft_worker.py: - Keep val_num_workers/val_prefetch_factor alongside new custom_val_data_loader flag - Keep pre_training_hook alongside new dp_load_balancing logic - Keep checkpoint skip/post_training_hook alongside new upload_thread wait
- Remove per-rank validation debug log block in sft_trainer.py - Change [DEBUG] to [SFT] prefix in epoch checkpoint log - Change [SFT Worker] to [SFT] in validation trigger log - Change [Validation] to [SFT] Validation in batch debug log
Resolve conflict in sft_worker.py: keep val_num_workers/val_prefetch_factor variables, drop unused custom_val_data_loader flag (removed upstream).
Resolve sft_trainer checkpointing: keep epoch-based naming and dtype export; match upstream rank-0 + barrier structure.
# Conflicts: # cosmos_rl/policy/config/__init__.py # cosmos_rl/policy/trainer/llm_trainer/sft_trainer.py
…hanced-hooks-and-custom-loggers' Rarunachalam/rc176 script option only See merge request nvidia-tao-toolkit/cosmos-reason!2
…hooks-and-custom-loggers' DAFT SFT hook See merge request nvidia-tao-toolkit/cosmos-reason!3
…s-and-custom-loggers # Conflicts: # pyproject.toml
7 participants
No description provided.