← Back to reference

Supported Framework Matrix

FrameworkVersionStatus
PyTorch FSDP2.1+Full
DeepSpeed0.13+Full
Megatron-LMcore_r0.6+Full
JAX0.4.20+Full
TensorFlow2.14+Full
Ray2.7+Full
KubeRay1.13+Full
PyTorch Lightning2.0+Best effort

Per-framework Notes

PyTorch

  • Forward / backward hooks via register_forward_hook / register_full_backward_hook.
  • FSDP shard metadata via summon_full_params.
  • DDP context captured.

DeepSpeed

  • Reads model_engine.communication_data_type.
  • Optimizer partition info.
  • ZeRO stage + offload.

Megatron-LM

  • TransformerLayer forward hooks.
  • TP/PP/CP/SP sizes + virtual pipeline model size.
  • Microbatch + pipeline bubble timing.

JAX

  • flax.linen.Module.sow tracing hook.
  • HLO via pjit_xla_computation.
  • Compile-time seconds.

TensorFlow

  • tf.profiler.experimental.start / stop.
  • tf.distribute.Strategy.experimental_distribute_dataset.
  • Checkpoint integrity validator.

Test Fixtures

denpex_sdk/tests/fixtures/ ships real traceback samples for torch OOM, FSDP all-gather OOM, DeepSpeed ZeRO-3 OOM, Megatron TP, JAX HLO, TF XLA, NCCL async error, NCCL watchdog.