Skip to content

feat: add TAO WTS and AETC SFT support - #160

Draft
ramanathan831 wants to merge 4 commits into
NVIDIA:mainfrom
ramanathan831:rarunachalam/tao-support-wts
Draft

feat: add TAO WTS and AETC SFT support#160
ramanathan831 wants to merge 4 commits into
NVIDIA:mainfrom
ramanathan831:rarunachalam/tao-support-wts

Conversation

@ramanathan831

Copy link
Copy Markdown

Summary

  • add TAO-compatible WTS supervised fine-tuning recipes for Cosmos Nano and Edge
  • add TAO status reporting, epoch-based validation/checkpoint integration, and non-root container support
  • add DAFT-backed AETC supervised fine-tuning configuration
  • exclude Python bytecode from Docker build contexts

Why

This brings the TAO training workflow from the Cosmos-RL integration into Cosmos Framework so fine-tuned Framework checkpoints can be trained and consumed by TAO actions without merging the public Cosmos-RL implementation into the runtime image.

Impact

Users can launch WTS and AETC SFT through Cosmos Framework with TAO-compatible reporting and checkpoint behavior. The Docker image can run under non-root Pyxis and SLURM environments.

Validation

  • built the Cosmos Framework container image
  • completed full WTS Nano and Edge training on SLURM
  • completed full AETC Nano and Edge training on the parity dataset
  • completed 367-sample AETC evaluations and compared Cosmos-RL and Framework BERTScore results

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