Skip to content

[ENH] Added TimeMoE forecaster to ptf-v2#2143

Open
lucifer4073 wants to merge 2 commits intosktime:mainfrom
lucifer4073:timemoe
Open

[ENH] Added TimeMoE forecaster to ptf-v2#2143
lucifer4073 wants to merge 2 commits intosktime:mainfrom
lucifer4073:timemoe

Conversation

@lucifer4073
Copy link
Contributor

Reference Issues/PRs

Closes #2137.

What does this implement/fix? Explain your changes.

Implementation adds the TimeMoE model in ptf-v2 inheriting TSLibBaseModel. Most of the implementation is similar to that in Thuml Time series library, except the following changes

  • training argument which if set to false freezes the model backbone and uses .generate() method for zero-shot-inference. Defaults to True and passes data directly to backbone.
  • used logits to get future_flat.

What should a reviewer concentrate their feedback on?

Implementation adds transformer dependency to ptf-v2.

Did you add any tests for the change?

pkg file consists of all the test args.

Any other comments?

PR checklist

  • The PR title starts with either [ENH], [MNT], [DOC], or [BUG]. [BUG] - bugfix, [MNT] - CI, test framework, [ENH] - adding or improving code, [DOC] - writing or improving documentation or docstrings.
  • Added/modified tests
  • Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
    To run hooks independent of commit, execute pre-commit run --all-files

@lucifer4073 lucifer4073 marked this pull request as ready for review March 4, 2026 21:09
@PranavBhatP
Copy link
Contributor

PranavBhatP commented Mar 5, 2026

@lucifer4073 are you sure about interfacing TimeMoE as a child class of TslibBaseModel? If I'm not wrong it's a foundational model and that would require an entirely different D2 layer, whose design hasn't been consolidated or discussed, please suggest a prior design before making large PRs. It gets very hard to review without prior context or discussion on the design. There is parallel work going on with other foundational models like Chronos for v2 you can catch up with that discussion, and see whether your current design is feasible

@codecov
Copy link

codecov bot commented Mar 5, 2026

Codecov Report

❌ Patch coverage is 91.76471% with 7 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@1952984). Learn more about missing BASE report.

Files with missing lines Patch % Lines
pytorch_forecasting/models/timemoe/_timemoe_v2.py 88.00% 6 Missing ⚠️
...asting/tests/test_all_v2/test_all_estimators_v2.py 90.90% 1 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2143   +/-   ##
=======================================
  Coverage        ?   86.67%           
=======================================
  Files           ?      168           
  Lines           ?     9817           
  Branches        ?        0           
=======================================
  Hits            ?     8509           
  Misses          ?     1308           
  Partials        ?        0           
Flag Coverage Δ
cpu 86.67% <91.76%> (?)
pytest 86.67% <91.76%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@lucifer4073
Copy link
Contributor Author

@lucifer4073 are you sure about interfacing TimeMoE as a child class of TslibBaseModel? If I'm not wrong it's a foundational model and that would require an entirely different D2 layer, whose design hasn't been consolidated or discussed, please suggest a prior design before making large PRs. It gets very hard to review without prior context or discussion on the design. There is parallel work going on with other foundational models like Chronos for v2 you can catch up with that discussion, and see whether your current design is feasible

@PranavBhatP Thanks for the comment. I was under the impression that all models in the THUML Time-Series module inherit from TslibBaseModel.
Regarding TimeMoE, it primarily leverages Hugging Face’s AutoModelForCausalLM for training and inference, and from what I can see, it aligns well with the existing TSlib design (with added dependency).
That said, I’ll review the Chronosv2 discussion more carefully to determine whether it would require a separate or entirely different D2 layer implementation. Thanks again for pointing this out.

@PranavBhatP
Copy link
Contributor

PranavBhatP commented Mar 5, 2026

I was under the impression that all models in the THUML Time-Series module inherit from TslibBaseModel.

Yes, you are right, but the "models" here are only trainable encoder decoder models, not pre-trained models with imported weights. There are already some models like TimeXer which are based on the transformer architecture, so you are not completely wrong here. I think your work on TimeMoe would definitely help find common patterns in creating a design for interfacing FMs into the library along with the work on Chronos.

@PranavBhatP
Copy link
Contributor

PranavBhatP commented Mar 5, 2026

Also just another comment, try sticking to a smaller number of PRs. It would help you prioritize efforts on quality. The emphasis is on Quality > Quantity. I see almost 4 active model interface PRs from your end, but the effort is appreciated. This is just an advice from my end :)

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.

[ENH] Add TimeMOE Model (THUML) to ptf-v2

2 participants