Skip to content

[FEATURE] Add profiling for FP8, MXFP8, and NVFP4 #1679

Description

@Shubham-Rasal

Miles supports several low-precision recipes, but right now it’s hard to see where the time
is actually going.

It would be useful to break a training step into:

  • Scale/amax calculation
  • Quantization
  • GEMM
  • Dequantization
  • Layout or memory operations

It would also help to collect basic stats such as saturation, underflow, NaN/Inf counts,
tensor shapes, and dtypes.

A small benchmark could compare BF16, FP8, MXFP8, and NVFP4 using the same model and batch
settings:

Recipe | Tokens/s | Step time | Scale time | Quant time | GEMM time
BF16 | ... | ... | N/A | N/A | ...
MXFP8 | ... | ... | ... | ... | ...
NVFP4 | ... | ... | ... | ... | ...

This could build on the existing profiler in miles/utils/profile_utils.py, using CUDA events
or NVTX ranges.

The profiler should be off by default, have almost no overhead when disabled, and not change
training results or checkpoint formats.

Related: #615

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions