Skip to content

Conversation

@SamirMoustafa
Copy link
Member

Add FLOPS estimation functionality

✨ Context

This PR adds FLOPS (Floating Point Operations) estimation capability to the ModelBase class, addressing issue #180 for model complexity/resource utilization comparison.

🧠 Rationale behind the change

FLOPS estimation is needed to compare model resource requirements beyond just parameter count, as parameter count doesn't linearly correlate with computational complexity.

Type of changes

  • ✨ New Feature (changes that introduce new functionality)
  • ✅ Tests (Unit tests, integration tests, end-to-end tests)

🛠 What does this PR implement

  • Add estimate_flops() method to ModelBase using torch.utils.flop_counter.FlopCounterMode
  • Add _resolve_method() helper to resolve method paths for FLOPS counting (supports nested methods like conch.forward)
  • Add MODEL_INPUT_ARGS dictionary with test inputs for all models in the registry
  • Add FLOPS estimation tests in test_models_general.py that verify FLOPS can be computed for each model

🙈 Missing

  • Documentation page with comparison table (parameter count, FLOPS, inference time) - to be added in follow-up PR
  • Inference time measurement - to be added in follow-up PR

🧪 How should this be tested?

  • uv run pytest tests/models/test_models_general.py::test_model_init - verifies FLOPS estimation works for all models

Addresses #180

- Add estimate_flops method to count FLOPS using torch.utils.flop_counter
- Add _resolve_method helper to resolve method paths for FLOPS counting
- Add MODEL_INPUT_ARGS dictionary with test inputs for all models
- Add FLOPS estimation tests in test_models_general.py

Addresses issue #180: Model complexity/resource utilization
@SamirMoustafa SamirMoustafa requested a review from Mr-Milk January 7, 2026 15:41
…dling

- Update estimate_flops method signature to accept method as the first argument followed by *args and **kwargs.
- Adjust test cases in test_models_general.py to align with the new method signature, ensuring proper FLOPS estimation functionality is maintained.
@Mr-Milk
Copy link
Member

Mr-Milk commented Jan 8, 2026

@SamirMoustafa I think we can safely merge this after the changes have been addressed. Could you please also update the model card with an additional FLOPS entry in a different PR?

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.

3 participants