Motivation
As Neural-LAM grows to support more diverse model architectures (Hi-LAM, Parallel) and complex evaluation workflows, our current configuration infrastructure is facing several structural challenges. This RFC aims to consolidate the configuration-related challenges that have emerged across the codebase and define the requirements for a more scalable and maintainable configuration architecture.
Identified Problems
Based on investigating the current configuration, logging, metrics, and checkpointing flows, I identified a few recurring areas of friction:
- Configuration Fragmentation: Configuration is currently split across YAML files and CLI arguments. As a result, some settings need to be manually connected and validated in multiple places.
- Unclear Module Dependencies: Some core modules still depend on broad configuration objects (for example, args: Any) rather than explicit inputs. This makes it harder to understand what a module actually needs and reduces ease of programmatic use.
- Validation Latency: Certain configuration errors are only detected after training or evaluation has already started. This can lead to failures that could have been caught earlier during setup.
- Configuration Logic Spread Across the Codebase: Configuration for logging, metrics, plotting, schedulers, and similar features is handled across multiple locations. This makes the behavior harder to follow and extend.
- Checkpoint and Configuration Coupling: Model checkpoints currently depend on parts of the configuration structure used when the model was created. This makes configuration refactoring more difficult while maintaining compatibility with older checkpoints.
Design Requirements & Constraints
Based on the discussion so far, any future configuration approach would ideally satisfy the following requirements:
1. Preserve Explicit APIs: Core classes should continue to expose clear and explicit constructor arguments rather than depending on large, catch-all configuration objects. This keeps dependencies visible and makes the code easier to understand and maintain.
2. Support Programmatic Usage: Neural-LAM should remain easy to use from notebooks and custom Python scripts, without requiring users to go through a specific CLI workflow.
3. Keep Configuration Loading Separate from Runtime Logic: Whether configuration comes from CLI arguments, YAML files, or another framework, runtime components should primarily work with validated values rather than being tightly coupled to a particular configuration system.
4. Maintain Backward Compatibility: Existing checkpoints and workflows should continue to work wherever practical, with migration paths available when changes are required.
Possible Directions
There are several possible ways to address the issues outlined above, ranging from improvements to the current approach to adopting a more structured configuration framework.
During discussion, a few candidate directions have already been mentioned, including type-driven approaches (e.g. Tyro), configuration frameworks (e.g. Hydra/OmegaConf), and further evolution of the current dataclass-based approach.
The approaches above are not intended to be exhaustive. I'm interested in hearing which directions seem most aligned with Neural-LAM's long-term goals and whether there are other alternatives worth considering.
I’m looking forward to the team's thoughts on these requirements and the best way to move forward!
Motivation
As Neural-LAM grows to support more diverse model architectures (Hi-LAM, Parallel) and complex evaluation workflows, our current configuration infrastructure is facing several structural challenges. This RFC aims to consolidate the configuration-related challenges that have emerged across the codebase and define the requirements for a more scalable and maintainable configuration architecture.
Identified Problems
Based on investigating the current configuration, logging, metrics, and checkpointing flows, I identified a few recurring areas of friction:
Design Requirements & Constraints
Based on the discussion so far, any future configuration approach would ideally satisfy the following requirements:
1. Preserve Explicit APIs: Core classes should continue to expose clear and explicit constructor arguments rather than depending on large, catch-all configuration objects. This keeps dependencies visible and makes the code easier to understand and maintain.
2. Support Programmatic Usage: Neural-LAM should remain easy to use from notebooks and custom Python scripts, without requiring users to go through a specific CLI workflow.
3. Keep Configuration Loading Separate from Runtime Logic: Whether configuration comes from CLI arguments, YAML files, or another framework, runtime components should primarily work with validated values rather than being tightly coupled to a particular configuration system.
4. Maintain Backward Compatibility: Existing checkpoints and workflows should continue to work wherever practical, with migration paths available when changes are required.
Possible Directions
There are several possible ways to address the issues outlined above, ranging from improvements to the current approach to adopting a more structured configuration framework.
During discussion, a few candidate directions have already been mentioned, including type-driven approaches (e.g. Tyro), configuration frameworks (e.g. Hydra/OmegaConf), and further evolution of the current dataclass-based approach.
The approaches above are not intended to be exhaustive. I'm interested in hearing which directions seem most aligned with Neural-LAM's long-term goals and whether there are other alternatives worth considering.
I’m looking forward to the team's thoughts on these requirements and the best way to move forward!