JET for evo2: 1b model training

[dorotat-nv]

Description

Adding 1b evo2 model training for partial conv with tflops logging. Step number to be adjusted later

Type of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor
• Documentation update
• Other (please describe):

CI Pipeline Configuration

Configure CI behavior by applying the relevant labels:

• SKIP_CI - Skip all continuous integration tests
• INCLUDE_NOTEBOOKS_TESTS - Execute notebook validation tests in pytest
• INCLUDE_SLOW_TESTS - Execute tests labelled as slow in pytest for extensive testing

Note

By default, the notebooks validation tests are skipped unless explicitly enabled.

Usage

TODO: Add code snippet

Pre-submit Checklist

• I have tested these changes locally
• I have updated the documentation accordingly
• I have added/updated tests as needed
• All existing tests pass successfully