TensorFlow two-worker mixed_bfloat16 Adam: AUTO collectives diverge with step scheduler and batch 2

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.18.0

Custom code

Yes

OS platform and distribution

Linux 5.15.0-113-generic x86_64 glibc 2.35

Mobile device

No response

Python version

3.11.15

Bazel version

No response

GCC/compiler version

14.3.0

CUDA/cuDNN version

12.5.1/9

GPU model and memory

Tesla V100S-PCIE-32GB / 32768 MiB

Current behavior?

The reference arm reaches a zero final loss in this replay, while the compiled arm stays at 13.85 with a nonzero gradient norm. The workers agree on those values, so the mismatch is not caused by one worker consuming different input.

Because the optimizer update is not the compiled portion, the result points toward the compiled forward/backward computation or the distributed gradient aggregation around it. The small batch and AUTO collective selection make this a compact reproducer for a distributed mixed-bfloat16 inconsistency.

Observed final metrics:

worker	reference loss	compiled loss	loss abs diff	loss tolerance	reference grad norm	compiled grad norm	grad abs diff	grad tolerance
worker 0	0	13.85247993	13.85247993	1.535247993	0	4.268177986	4.268177986	1.853635597
worker 1	0	13.85247993	13.85247993	1.535247993	0	4.268177986	4.268177986	1.853635597

The JSON config attached to the report is the small-batch Adam case used here; it specifies AUTO communication, mixed_bfloat16, clipping, and the two worker addresses.

Expected behavior

The compiled path should not prevent the same two-worker Adam step from reaching the same zero-loss state observed by the reference path. Some bfloat16 noise is expected, but a nonzero-loss branch at this scale should remain inside tolerance if the distributed gradients are equivalent.

Standalone code to reproduce the issue

Machine 1 / worker 1:


python minimal_repro_mwms_xla.py \
  --config tensorflow_official_repro_config.json \
  --worker-index 1 \
  --worker-addresses <machine0_ip>:61120,<machine1_ip>:62120 \
  --cuda-visible-devices 0


Machine 0 / worker 0:


python minimal_repro_mwms_xla.py \
  --config tensorflow_official_repro_config.json \
  --worker-index 0 \
  --worker-addresses <machine0_ip>:61120,<machine1_ip>:62120 \
  --cuda-visible-devices 0


minimal_repro_mwms_xla.py

#!/usr/bin/env python3
"""Minimal two-worker TensorFlow reproducer for MWMS XLA/non-XLA divergence.

Run this same file on two GPU machines, one process per machine.  The script is
standalone TensorFlow code: it uses deterministic synthetic tensors and does not
depend on the original project pipeline or an external dataset.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any


LOSS_ATOL = 0.15
LOSS_RTOL = 0.10
GRAD_ATOL = 1.0
GRAD_RTOL = 0.20


class TeeOutput:
    def __init__(self, filepath: Path):
        self.file = filepath.open("w", buffering=1, encoding="utf-8")
        self.stdout = sys.stdout

    def write(self, message: str) -> None:
        self.stdout.write(message)
        self.file.write(message)

    def flush(self) -> None:
        self.stdout.flush()
        self.file.flush()

    def close(self) -> None:
        self.file.close()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=None, help="JSON runtime configuration.")
    parser.add_argument("--worker-index", type=int, required=True)
    parser.add_argument("--worker-addresses", required=True, help="host0:port0,host1:port1")
    parser.add_argument("--output-dir", default=None)
    parser.add_argument("--cuda-visible-devices", default="0")
    parser.add_argument("--steps", type=int, default=20)
    parser.add_argument("--per-replica-batch-size", type=int, default=4)
    parser.add_argument("--input-size", type=int, default=224)
    parser.add_argument("--num-classes", type=int, default=10)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--learning-rate", type=float, default=0.001)
    parser.add_argument("--weight-decay", type=float, default=0.0)
    parser.add_argument("--optimizer", choices=["adam", "adamw", "sgd"], default="adam")
    parser.add_argument("--momentum", type=float, default=0.0)
    parser.add_argument("--precision-policy", default="float32")
    parser.add_argument("--communication", choices=["RING", "AUTO"], default="RING")
    parser.add_argument("--intra-op-threads", type=int, default=2)
    parser.add_argument("--inter-op-threads", type=int, default=2)
    parser.add_argument("--fail-on-divergence", action="store_true")
    return parser.parse_args()


def setup_output_directory(args: argparse.Namespace) -> tuple[Path, TeeOutput]:
    if args.output_dir:
        output_dir = Path(args.output_dir)
    else:
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        config_name = Path(args.config).stem if args.config else "config"
        output_dir = Path(f"minimal_repro_{config_name}_{args.steps}step_run_{timestamp}")
    output_dir.mkdir(parents=True, exist_ok=True)
    tee = TeeOutput(output_dir / f"worker{args.worker_index}.stdout")
    (output_dir / f"worker{args.worker_index}_started.txt").write_text(
        f"Started at {datetime.now()}\n", encoding="utf-8"
    )
    return output_dir, tee


def apply_config(args: argparse.Namespace) -> dict[str, Any]:
    if not args.config:
        return {}
    path = Path(args.config).expanduser().resolve()
    cfg = json.loads(path.read_text(encoding="utf-8"))
    args.config = str(path)

    args.num_classes = int(cfg.get("num_classes", args.num_classes))
    args.seed = int(cfg.get("seed", args.seed))
    args.learning_rate = float(cfg.get("learning_rate", args.learning_rate))
    args.weight_decay = float(cfg.get("weight_decay", args.weight_decay))
    args.optimizer = str(cfg.get("optimizer", args.optimizer)).lower()
    args.momentum = float(cfg.get("momentum", args.momentum))
    args.precision_policy = str(cfg.get("precision_policy", args.precision_policy))
    args.intra_op_threads = int(cfg.get("intra_op_parallelism_threads") or args.intra_op_threads)
    args.inter_op_threads = int(cfg.get("inter_op_parallelism_threads") or args.inter_op_threads)

    shape = cfg.get("input_shape_per_replica")
    if isinstance(shape, list) and len(shape) >= 3:
        args.per_replica_batch_size = int(shape[0])
        args.input_size = int(shape[1])
    else:
        args.per_replica_batch_size = int(cfg.get("per_replica_batch_size", args.per_replica_batch_size))
    if "steps_per_arm" in cfg:
        args.steps = int(cfg["steps_per_arm"])

    communication = str(
        cfg.get("communication")
        or cfg.get("communication_implementation")
        or args.communication
    ).upper()
    if communication in {"RING", "AUTO"}:
        args.communication = communication
    return cfg


def setup_environment(args: argparse.Namespace) -> list[str]:
    workers = [part.strip() for part in args.worker_addresses.split(",") if part.strip()]
    if len(workers) != 2:
        raise SystemExit("--worker-addresses must contain exactly two host:port entries")
    if args.worker_index not in (0, 1):
        raise SystemExit("--worker-index must be 0 or 1")
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_visible_devices)
    os.environ.setdefault("TF_FORCE_GPU_ALLOW_GROWTH", "true")
    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "1")
    os.environ["TF_CONFIG"] = json.dumps(
        {"cluster": {"worker": workers}, "task": {"type": "worker", "index": int(args.worker_index)}},
        sort_keys=True,
    )
    return workers


def config_summary(args: argparse.Namespace, cfg: dict[str, Any]) -> dict[str, Any]:
    if not cfg:
        return {}
    summary = {
        "config_file": Path(str(args.config)).name,
        "device": cfg.get("device"),
        "distributed_strategy": cfg.get("distributed_strategy"),
        "tf_strategy_class": cfg.get("tf_strategy_class"),
        "communication_implementation": cfg.get("communication_implementation") or cfg.get("communication"),
        "precision_policy": cfg.get("precision_policy"),
        "optimizer": cfg.get("optimizer"),
        "seed": cfg.get("seed"),
        "per_replica_batch_size": cfg.get("per_replica_batch_size"),
        "input_shape_per_replica": cfg.get("input_shape_per_replica"),
        "learning_rate": cfg.get("learning_rate"),
        "weight_decay": cfg.get("weight_decay"),
        "momentum": cfg.get("momentum"),
        "steps_per_arm": cfg.get("steps_per_arm"),
        "run_eagerly": cfg.get("run_eagerly"),
    }
    source_config = cfg.get("source_config_full")
    if isinstance(source_config, dict):
        summary["config_top_level_key_count"] = len(cfg)
        summary["source_config_full_key_count"] = len(source_config)
        summary["runtime_projection_overrides"] = cfg.get("runtime_projection_overrides", {})
    return summary


def runtime_configuration(args: argparse.Namespace, workers: list[str]) -> dict[str, Any]:
    return {
        "strategy": "MultiWorkerMirroredStrategy",
        "workers": workers,
        "num_workers": len(workers),
        "gpus_per_worker": 1,
        "communication": str(args.communication),
        "model": "ResNet50(weights=None, include_top=False, pooling='avg') + Dense(10)",
        "input": "synthetic normal tensor generated from a fixed seed",
        "steps_per_arm": int(args.steps),
        "per_replica_batch_size": int(args.per_replica_batch_size),
        "effective_global_batch_size": int(args.per_replica_batch_size) * len(workers),
        "input_shape_per_replica": [
            int(args.per_replica_batch_size),
            int(args.input_size),
            int(args.input_size),
            3,
        ],
        "num_classes": int(args.num_classes),
        "seed": int(args.seed),
        "optimizer": str(args.optimizer),
        "learning_rate": float(args.learning_rate),
        "weight_decay": float(args.weight_decay),
        "momentum": float(args.momentum),
        "precision_policy": str(args.precision_policy),
        "threading": {
            "intra_op_parallelism_threads": int(args.intra_op_threads),
            "inter_op_parallelism_threads": int(args.inter_op_threads),
        },
        "arms": [
            {"name": "nonxla", "forward_backward_jit_compile": False},
            {"name": "xla", "forward_backward_jit_compile": True},
        ],
        "optimizer_step_jit_compile": False,
        "comparison_thresholds": {
            "loss_atol": LOSS_ATOL,
            "loss_rtol": LOSS_RTOL,
            "grad_atol": GRAD_ATOL,
            "grad_rtol": GRAD_RTOL,
        },
    }


def make_model(tf: Any, input_size: int, num_classes: int) -> Any:
    base = tf.keras.applications.ResNet50(
        weights=None,
        include_top=False,
        input_shape=(input_size, input_size, 3),
        pooling="avg",
    )
    out = tf.keras.layers.Dense(num_classes)(base.output)
    return tf.keras.Model(base.input, out)


def make_optimizer(tf: Any, args: argparse.Namespace) -> Any:
    if args.optimizer == "sgd":
        try:
            return tf.keras.optimizers.SGD(
                learning_rate=float(args.learning_rate),
                momentum=float(args.momentum),
                weight_decay=float(args.weight_decay),
            )
        except TypeError:
            return tf.keras.optimizers.SGD(
                learning_rate=float(args.learning_rate),
                momentum=float(args.momentum),
            )
    if args.optimizer == "adamw":
        return tf.keras.optimizers.AdamW(
            learning_rate=float(args.learning_rate),
            weight_decay=float(args.weight_decay),
        )
    try:
        return tf.keras.optimizers.Adam(
            learning_rate=float(args.learning_rate),
            weight_decay=float(args.weight_decay),
        )
    except TypeError:
        return tf.keras.optimizers.Adam(learning_rate=float(args.learning_rate))


def run_arm(tf: Any, strategy: Any, args: argparse.Namespace, *, jit_compile: bool) -> dict[str, float]:
    tf.keras.utils.set_random_seed(int(args.seed))
    tf.keras.mixed_precision.set_global_policy(str(args.precision_policy))

    with strategy.scope():
        model = make_model(tf, int(args.input_size), int(args.num_classes))
        optimizer = make_optimizer(tf, args)
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True,
            reduction=tf.keras.losses.Reduction.NONE,
        )

    def grad_fn(x, y):
        with tf.GradientTape() as tape:
            logits = model(x, training=True)
            per_example_loss = loss_fn(y, logits)
            loss = tf.nn.compute_average_loss(per_example_loss)
        grads = tape.gradient(loss, model.trainable_variables)
        grad_norm = tf.linalg.global_norm([g for g in grads if g is not None])
        return loss, grads, grad_norm

    compiled_grad_fn = tf.function(grad_fn, jit_compile=bool(jit_compile))

    def step_fn(x, y):
        loss, grads, grad_norm = compiled_grad_fn(x, y)
        optimizer.apply_gradients(
            (g, v) for g, v in zip(grads, model.trainable_variables) if g is not None
        )
        return loss, grad_norm

    final_loss = float("nan")
    final_grad_norm = float("nan")
    batch = int(args.per_replica_batch_size)
    input_size = int(args.input_size)
    for step in range(int(args.steps)):
        generator = tf.random.Generator.from_seed(int(args.seed) + step)
        x = generator.normal((batch, input_size, input_size, 3))
        y = tf.range(batch, dtype=tf.int32) % int(args.num_classes)
        per_replica_loss, per_replica_grad_norm = strategy.run(step_fn, args=(x, y))
        final_loss = float(strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, axis=None).numpy())
        final_grad_norm = float(
            strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_grad_norm, axis=None).numpy()
        )
    return {"loss": final_loss, "grad_norm": final_grad_norm}


def compare(nonxla: dict[str, float], xla: dict[str, float]) -> dict[str, Any]:
    loss_abs = abs(nonxla["loss"] - xla["loss"])
    grad_abs = abs(nonxla["grad_norm"] - xla["grad_norm"])
    loss_tol = LOSS_ATOL + LOSS_RTOL * max(abs(nonxla["loss"]), abs(xla["loss"]))
    grad_tol = GRAD_ATOL + GRAD_RTOL * max(abs(nonxla["grad_norm"]), abs(xla["grad_norm"]))
    loss_div = math.isfinite(loss_abs) and loss_abs > loss_tol
    grad_div = math.isfinite(grad_abs) and grad_abs > grad_tol
    return {
        "divergent": bool(loss_div or grad_div),
        "nonxla_loss": nonxla["loss"],
        "xla_loss": xla["loss"],
        "loss_abs_diff": loss_abs,
        "loss_tolerance": loss_tol,
        "loss_divergent": bool(loss_div),
        "nonxla_grad_norm": nonxla["grad_norm"],
        "xla_grad_norm": xla["grad_norm"],
        "grad_abs_diff": grad_abs,
        "grad_tolerance": grad_tol,
        "grad_divergent": bool(grad_div),
    }


def main() -> int:
    args = parse_args()
    output_dir, tee = setup_output_directory(args)
    sys.stdout = tee
    print(f"# Output directory: {output_dir}", flush=True)
    print(f"# Worker index: {args.worker_index}", flush=True)
    print(f"# Steps: {args.steps}", flush=True)
    print("", flush=True)

    source_config = apply_config(args)
    workers = setup_environment(args)

    import tensorflow as tf

    try:
        tf.config.threading.set_intra_op_parallelism_threads(int(args.intra_op_threads))
        tf.config.threading.set_inter_op_parallelism_threads(int(args.inter_op_threads))
    except RuntimeError:
        pass

    gpus = tf.config.list_physical_devices("GPU")
    if not gpus:
        raise RuntimeError("This reproduction expects one visible GPU per worker.")
    for gpu in gpus:
        try:
            tf.config.experimental.set_memory_growth(gpu, True)
        except RuntimeError:
            pass

    communication_impl = (
        tf.distribute.experimental.CommunicationImplementation.RING
        if str(args.communication).upper() == "RING"
        else tf.distribute.experimental.CommunicationImplementation.AUTO
    )
    strategy = tf.distribute.MultiWorkerMirroredStrategy(
        communication_options=tf.distribute.experimental.CommunicationOptions(
            implementation=communication_impl
        )
    )

    env_row = {
        "event": "environment",
        "worker_index": int(args.worker_index),
        "tensorflow_version": tf.__version__,
        "gpus": [str(gpu) for gpu in gpus],
        "tf_config": json.loads(os.environ["TF_CONFIG"]),
        "workers": workers,
        "runtime_configuration": runtime_configuration(args, workers),
        "source_config": config_summary(args, source_config),
    }
    try:
        env_row["build_info"] = dict(tf.sysconfig.get_build_info())
    except Exception:
        env_row["build_info"] = {}
    print(json.dumps(env_row, sort_keys=True), flush=True)

    started = time.perf_counter()
    nonxla = run_arm(tf, strategy, args, jit_compile=False)
    print(json.dumps({"event": "arm", "worker_index": int(args.worker_index), "arm": "nonxla", **nonxla}, sort_keys=True), flush=True)

    xla = run_arm(tf, strategy, args, jit_compile=True)
    print(json.dumps({"event": "arm", "worker_index": int(args.worker_index), "arm": "xla", **xla}, sort_keys=True), flush=True)

    result = compare(nonxla, xla)
    print(
        json.dumps(
            {
                "event": "summary",
                "worker_index": int(args.worker_index),
                "status": "DIVERGENCE_REPRODUCED" if result["divergent"] else "not_reproduced",
                "wall_time_sec": round(time.perf_counter() - started, 3),
                **result,
            },
            sort_keys=True,
        ),
        flush=True,
    )

    if args.worker_index == 0:
        (output_dir / "returncodes.txt").write_text("0\n", encoding="utf-8")
    tee.close()
    sys.stdout = tee.stdout
    print(f"\nWorker {args.worker_index} completed. Output saved to: {output_dir}")
    if args.fail_on_divergence and result["divergent"]:
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())


tensorflow_official_repro_config.json

{
  "arms": [
    {
      "forward_backward_jit_compile": false,
      "name": "nonxla"
    },
    {
      "forward_backward_jit_compile": true,
      "name": "xla"
    }
  ],
  "auto_shard_policy": "AUTO",
  "backup_and_restore": false,
  "batch_size": 2,
  "boundary_target": "large_lr",
  "bytes_per_pack": null,
  "cache": true,
  "cache_file": null,
  "checkpoint_dir": "runs/tf_checkpoints",
  "checkpoint_enabled": false,
  "checkpoint_interval": 100,
  "clip_grad_norm": 5.0,
  "collective_timeout_sec": null,
  "communication_implementation": "AUTO",
  "comparison_thresholds": {
    "grad_atol": 1.0,
    "grad_rtol": 0.2,
    "loss_atol": 0.15,
    "loss_rtol": 0.1
  },
  "cross_device_ops": null,
  "data_service_enabled": false,
  "dataset": "cifar10",
  "deterministic_ops": false,
  "device": "gpu",
  "distributed_strategy": "multi_worker_mirrored",
  "drop_remainder": true,
  "effective_global_batch_size": 4,
  "framework": "tensorflow",
  "global_batch_size": 2,
  "gpu_memory_growth": true,
  "gpu_memory_limit_mb": 4096,
  "gradient_accumulation_steps": 1,
  "id": "tf_cf_boundary-00274",
  "input": "deterministic synthetic normal tensor",
  "input_options_fetch_to_device": true,
  "input_shape_per_replica": [
    2,
    224,
    224,
    3
  ],
  "inter_op_parallelism_threads": null,
  "interleave_cycle_length": 4,
  "intra_op_parallelism_threads": 2,
  "jit_compile": false,
  "learning_rate": 1.0,
  "local_device_count": 1,
  "log_interval": 5,
  "loss_divergence_threshold": 2.0,
  "loss_scale": "none",
  "lr_scaling": "none",
  "map_deterministic": false,
  "max_steps": 5,
  "model": "resnet50",
  "momentum": 0.9,
  "monitor_grad_norm": true,
  "monitor_memory": true,
  "monitor_step_time": true,
  "nan_inf_check": true,
  "num_classes": 10,
  "num_parallel_calls": "AUTOTUNE",
  "num_replicas_in_sync": 2,
  "num_workers": 2,
  "optimizer": "adam",
  "optimizer_step_jit_compile": false,
  "per_replica_batch_size": 2,
  "performance_regression_threshold": 0.3,
  "precision_policy": "mixed_bfloat16",
  "prefetch": "AUTOTUNE",
  "profile": "cf_boundary",
  "reference_global_batch_size": 256,
  "run_eagerly": false,
  "runtime_projection_overrides": {
    "arms": {
      "from": null,
      "to": [
        {
          "forward_backward_jit_compile": false,
          "name": "nonxla"
        },
        {
          "forward_backward_jit_compile": true,
          "name": "xla"
        }
      ]
    },
    "comparison_thresholds": {
      "from": null,
      "to": {
        "grad_atol": 1.0,
        "grad_rtol": 0.2,
        "loss_atol": 0.15,
        "loss_rtol": 0.1
      }
    },
    "distributed_strategy": {
      "from": "mirrored",
      "to": "multi_worker_mirrored"
    },
    "effective_global_batch_size": {
      "from": null,
      "to": 4
    },
    "input": {
      "from": null,
      "to": "deterministic synthetic normal tensor"
    },
    "input_shape_per_replica": {
      "from": null,
      "to": [
        2,
        224,
        224,
        3
      ]
    },
    "local_device_count": {
      "from": "auto",
      "to": 1
    },
    "num_replicas_in_sync": {
      "from": 1,
      "to": 2
    },
    "num_workers": {
      "from": 1,
      "to": 2
    },
    "optimizer_step_jit_compile": {
      "from": null,
      "to": false
    },
    "steps_per_arm": {
      "from": null,
      "to": 20
    },
    "tf_strategy_class": {
      "from": "MirroredStrategy",
      "to": "MultiWorkerMirroredStrategy"
    },
    "visible_devices": {
      "from": null,
      "to": "0"
    },
    "worker_addresses": {
      "from": [
        "127.0.0.1:41685"
      ],
      "to": [
        "10.60.88.171:61120",
        "10.60.210.23:62120"
      ]
    }
  },
  "runtime_projection_reason": "Standalone two-worker TensorFlow reproducer: keep original tunables where they do not prevent execution, but force MultiWorkerMirroredStrategy, two workers, deterministic synthetic input, one visible GPU per worker, and 20-step XLA/non-XLA arms.",
  "scheduler": "step",
  "seed": 0,
  "shuffle": true,
  "shuffle_buffer_size": 10000,
  "source_config_full": {
    "auto_shard_policy": "AUTO",
    "backup_and_restore": false,
    "batch_size": 2,
    "boundary_target": "large_lr",
    "bytes_per_pack": null,
    "cache": true,
    "cache_file": null,
    "checkpoint_dir": "runs/tf_checkpoints",
    "checkpoint_enabled": false,
    "checkpoint_interval": 100,
    "clip_grad_norm": 5.0,
    "collective_timeout_sec": null,
    "communication_implementation": "AUTO",
    "cross_device_ops": null,
    "data_service_enabled": false,
    "dataset": "cifar10",
    "deterministic_ops": false,
    "device": "gpu",
    "distributed_strategy": "mirrored",
    "drop_remainder": true,
    "framework": "tensorflow",
    "global_batch_size": 2,
    "gpu_memory_growth": true,
    "gpu_memory_limit_mb": 4096,
    "gradient_accumulation_steps": 1,
    "id": "tf_cf_boundary-00274",
    "input_options_fetch_to_device": true,
    "inter_op_parallelism_threads": null,
    "interleave_cycle_length": 4,
    "intra_op_parallelism_threads": 2,
    "jit_compile": false,
    "learning_rate": 1.0,
    "local_device_count": "auto",
    "log_interval": 5,
    "loss_divergence_threshold": 2.0,
    "loss_scale": "none",
    "lr_scaling": "none",
    "map_deterministic": false,
    "max_steps": 5,
    "model": "resnet50",
    "momentum": 0.9,
    "monitor_grad_norm": true,
    "monitor_memory": true,
    "monitor_step_time": true,
    "nan_inf_check": true,
    "num_classes": 10,
    "num_parallel_calls": "AUTOTUNE",
    "num_replicas_in_sync": 1,
    "num_workers": 1,
    "optimizer": "adam",
    "per_replica_batch_size": 2,
    "performance_regression_threshold": 0.3,
    "precision_policy": "mixed_bfloat16",
    "prefetch": "AUTOTUNE",
    "profile": "cf_boundary",
    "reference_global_batch_size": 256,
    "run_eagerly": false,
    "scheduler": "step",
    "seed": 0,
    "shuffle": true,
    "shuffle_buffer_size": 10000,
    "tf_config": null,
    "tf_strategy_class": "MirroredStrategy",
    "tf_task_index": 0,
    "tf_task_type": "worker",
    "throughput_baseline": null,
    "visible_devices": null,
    "warmup_steps": 0,
    "weight_decay": 0.0,
    "worker_addresses": [
      "127.0.0.1:41685"
    ]
  },
  "steps_per_arm": 20,
  "tf_config": null,
  "tf_strategy_class": "MultiWorkerMirroredStrategy",
  "tf_task_index": 0,
  "tf_task_type": "worker",
  "throughput_baseline": null,
  "visible_devices": "0",
  "warmup_steps": 0,
  "weight_decay": 0.0,
  "worker_addresses": [
    "10.60.88.171:61120",
    "10.60.210.23:62120"
  ]
}

Relevant log output

{"divergent": true, "event": "summary", "grad_abs_diff": 4.2681779861450195, "grad_divergent": true, "grad_tolerance": 1.853635597229004, "loss_abs_diff": 13.852479934692383, "loss_divergent": true, "loss_tolerance": 1.5352479934692382, "nonxla_grad_norm": 0.0, "nonxla_loss": 0.0, "status": "DIVERGENCE_REPRODUCED", "wall_time_sec": 98.065, "worker_index": 0, "xla_grad_norm": 4.2681779861450195, "xla_loss": 13.852479934692383}
{"divergent": true, "event": "summary", "grad_abs_diff": 4.2681779861450195, "grad_divergent": true, "grad_tolerance": 1.853635597229004, "loss_abs_diff": 13.852479934692383, "loss_divergent": true, "loss_tolerance": 1.5352479934692382, "nonxla_grad_norm": 0.0, "nonxla_loss": 0.0, "status": "DIVERGENCE_REPRODUCED", "wall_time_sec": 98.066, "worker_index": 1, "xla_grad_norm": 4.2681779861450195, "xla_loss": 13.852479934692383}