Skip to content

Dev rocm mx type - #94

Closed
ScXfjiang wants to merge 1 commit into
rocm-mainfrom
dev_rocm_mx_type
Closed

ScXfjiang wants to merge 1 commit into
rocm-mainfrom
dev_rocm_mx_type

Conversation

@ScXfjiang

Copy link
Copy Markdown

No description provided.

quantization=jax_mxfp8 / jax_mxfp4 runs the dense projections, forward and both backward
GEMMs, through jax.nn.scaled_dot_general with one OCP MX config: E4M3 / E2M1 elements with
an E8M0 scale per block of 32. The recipe is stateless, so checkpoints are those of the
unquantized model.

Configs with use_qwix_quantization=True, shard_mode=explicit or MoE models are rejected.
jax_mxfp4 needs a JAX with the mxfp4 config (jax-ml/jax#41341).
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.

1 participant