[JAX] Relax tolerance for the test_multiprocessing_encoder.py with NVFP4 by 0.001

[phu0ngng]

Description

Relax tolerance for the test_multiprocessing_encoder.py with NVFP4 by 0.001

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

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 Overview

Greptile Summary

Relaxed the training accuracy tolerance threshold for NVFP4 (4-bit floating point) tests from 0.788 to 0.787 (0.001 decrease) in test_te_nvfp4() and test_te_nvfp4_shardy().

• NVFP4 is a low-precision quantization format that trades accuracy for performance
• The change affects two test methods that validate NVFP4BlockScaling recipe with multi-GPU multiprocessing
• Original tolerance (0.788) was set when NVFP4 support was added in PR [JAX] NVFP4 support in TE/JAX #2254
• The 0.1% relaxation is minimal and consistent with the inherent accuracy variance in low-precision quantization
• No changes to loss tolerance (< 0.451) were made

Confidence Score: 4/5

• This PR is safe to merge with low risk
• The change is minimal and well-scoped: relaxing a test accuracy tolerance by 0.001 (0.1%) for NVFP4 low-precision quantization tests. However, lacking documentation (commit message says only "relax tol") and no evidence provided (e.g., CI logs, empirical test results) to justify why this specific relaxation is needed reduces confidence from 5 to 4
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
examples/jax/encoder/test_multiprocessing_encoder.py	4/5	Relaxed accuracy tolerance from 0.788 to 0.787 for NVFP4 tests (both standard and shardy variants)

Sequence Diagram

sequenceDiagram
participant Test as test_te_nvfp4/test_te_nvfp4_shardy
participant Exec as exec()
participant Train as train_and_evaluate()
participant Assert as Assertion Check
Test->>Exec: Call with NVFP4BlockScaling recipe
Exec->>Train: Run 5 epochs with NVFP4 quantization
Train-->>Exec: Return [train_loss, train_accuracy, test_loss, test_accuracy]
Exec-->>Test: Return result
Test->>Assert: Check result[0] < 0.451 and result[1] > 0.787
Note over Assert: OLD: result[1] > 0.788<br/>NEW: result[1] > 0.787<br/>(tolerance relaxed by 0.001)

Loading

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[jberchtold-nvidia]

LGTM, thanks!

[phu0ngng]

/te-ci JAX L0