[JAX] Fix test_layer to support fused attention and adjust test encoder tolerance to account for minor diff

[jberchtold-nvidia]

Description

Fixes failing CI in test_layer.py and test_model_parallel_encoder.py

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Changed test_layer.py to map layer names that can support both unfused and fused attention layer names compared to a ref pure-JAX impl
• Slightly relaxed tolerances of delayedscaling test encoder from loss<0.361 to 0.362

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR fixes failing CI tests in the JAX test suite by addressing two distinct issues:

• test_layer.py: Added parameter name transformations for _FusedDotProductAttention_0/softmax_offset to complement existing _UnfusedDotProductAttention_0 mappings. The test compares TransformerEngine layer implementations against reference pure-JAX implementations by syncing parameters between them. When fused attention is enabled (determined by environment or runtime), the parameter paths include _FusedDotProductAttention_0 instead of _UnfusedDotProductAttention_0, causing the parameter sync to fail without these mappings.

• test_model_parallel_encoder.py: Adjusted loss tolerance threshold from < 0.361 to < 0.362 in four DelayedScaling FP8 tests. This accommodates minor numerical differences in FP8 training that can cause loss values to slightly exceed the previous threshold while still indicating correct behavior.

Confidence Score: 4/5

• This PR is safe to merge with low risk - fixes straightforward test infrastructure issues
• The changes are minimal and well-scoped test fixes. The test_layer.py change adds missing parameter mappings following an established pattern, ensuring compatibility with both fused and unfused attention implementations. The tolerance adjustment is conservative (0.001 increase) and only affects DelayedScaling FP8 tests where minor numerical variance is expected. However, the tolerance change deserves verification that actual loss values are within expected bounds
• Verify that the actual failing loss values in CI were close to 0.361-0.362 range to confirm the tolerance adjustment is appropriate

Important Files Changed

Filename	Overview
tests/jax/test_layer.py	Added fused attention layer name mappings to support both _FusedDotProductAttention_0 and _UnfusedDotProductAttention_0 paths in param transformations
examples/jax/encoder/test_model_parallel_encoder.py	Relaxed loss tolerance from 0.361 to 0.362 in DelayedScaling FP8 tests to account for numerical differences

Sequence Diagram

sequenceDiagram
    participant Test as Test Runner
    participant TLayer as TransformerEngine Layer
    participant RefLayer as Reference JAX Layer
    participant Sync as Parameter Sync
    participant Attn as Attention Module
    
    Note over Test,Attn: test_layer.py - Parameter Mapping Fix
    Test->>TLayer: Initialize with config
    Test->>RefLayer: Initialize reference layer
    TLayer->>Attn: Create DotProductAttention
    alt Fused Attention Enabled
        Attn-->>TLayer: _FusedDotProductAttention_0/softmax_offset
    else Unfused Attention
        Attn-->>TLayer: _UnfusedDotProductAttention_0/softmax_offset
    end
    Test->>Sync: sync_params_values(target, ref, transformations)
    Sync->>Sync: Map TE params to Ref params
    Note over Sync: Now supports both:<br/>_FusedDotProductAttention_0<br/>_UnfusedDotProductAttention_0
    Sync-->>Test: Synchronized parameters
    Test->>TLayer: Forward pass with synced params
    Test->>RefLayer: Forward pass
    Test->>Test: Compare outputs
    
    Note over Test: test_model_parallel_encoder.py - Tolerance Adjustment
    Test->>Test: train_and_evaluate() with FP8
    loop Training epochs
        Test->>Test: Train with DelayedScaling FP8
    end
    Test->>Test: Calculate final loss
    alt Loss within tolerance
        Test->>Test: assert loss < 0.362 ✓
    else Loss exceeds old tolerance
        Test->>Test: Would fail at 0.361, passes at 0.362
    end

Loading

[jberchtold-nvidia]

/te-ci L0 L1 L2 jax

[tdophung]

LGTM