zhangqiaorjc

zhangqiaorjc

zhangqiaorjc

import numpy as np
import os, re
import jax
from jax.experimental import maps
from jax.experimental import pjit
import jax.numpy as jnp

from jax.experimental import mesh_utils

from absl import flags

zhangqiaorjc

"""Decoder-only LM scaling experiments on GPUs."""

from jax import numpy as jnp
from paxml import experiment_registry
from paxml.tasks.lm.params.lm_cloud import LmCloudSpmd
from paxml.tasks.lm.params.lm_cloud import LmCloudSpmdPipeline
from praxis import layers


# TODO(zhangqiaorjc): Might need to use pmap instead of pjit for smaller models.

zhangqiaorjc

zhangqiaorjc

zhangqiaorjc

zhangqiaorjc

"""Runs a simple mnist model with fake FP8. FP8 scaling is used.

The HLO can be dumped by setting the environment variable:

  XLA_FLAGS='--xla_dump_disable_metadata=true --xla_dump_to=/tmp/hlo'


"""
import tensorflow as tf

zhangqiaorjc

from absl.testing import absltest
from absl import logging
import jax
import jax.numpy as jnp


def amax(x):
  return jnp.max(jnp.abs(x))

zhangqiaorjc

def make_hlo(f, optimize=False, metadata=False, platform=None):
  """Utility function for printing JAX-emitted HLO and XLA-compiled HLO.

  Args:
    f: jax function to return hlo for.
    optimize: bool: whether to return platform-specific, XLA-optimized HLO
    metadata: bool: whether to include JAX metadata information
    platform: Optional[str]: None, 'cpu','gpu','tpu' - platform to compile for,
      None uses default.