eager vs tf.function(jit_compile=False): forward matches to 5e-7 but a single backward step diverges on 161/214 gradient tensors, gated by Grappler (ResNet50, CPU)

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

binary

TensorFlow version

2.21.0 (Keras 3.15.1)

Custom code

Yes

OS platform and distribution

Ubuntu 22.04.1 LTS, Linux x86_64, kernel 5.15.0-78-generic, Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz (80 logical CPUs: 2 sockets x 20 cores x 2 threads). CPU-only run: no CUDA driver is visible (cuInit reports CUDA_ERROR_NO_DEVICE) and TensorFlow sees a single logical device /device:CPU:0. Intra/inter-op parallelism left at default; oneDNN custom operations left at their default (enabled).

Mobile device

No response

Python version

3.11.16

Bazel version

No response (installed from a binary wheel, not built from source)

GCC/compiler version

No response (installed from a binary wheel; the host system compiler is gcc (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0, but the compiler used to build the wheel was not independently established)

CUDA/cuDNN version

No response (CPU-only run; no CUDA-capable device is detected)

GPU model and memory

No response (no GPU visible to TensorFlow)

Current behavior?

On a single CPU device I compared the backward gradients of eager execution against tf.function(jit_compile=False) ("graph") for ONE forward/backward step of tf.keras.applications.ResNet50 (weights=None, random init) with BatchNorm in training mode, batch size 2, 224x224x3 synthetic input (tf.random.stateless_normal). Both arms start from bit-identical serialised initial weights and each runs in its own child process. I ran the comparison twice: once with Grappler's meta-optimizer at its default (on) and once with it disabled (tf.config.optimizer.set_experimental_options({"disable_meta_optimizer": True})).

The forward pass agrees; the backward pass does not, and the backward divergence is gated by Grappler:

	Grappler ON	Grappler OFF
forward loss rel. diff (eager vs graph)	5.3e-07	0.0 (bitwise)
backward tensors over tolerance	161 / 214	1 / 214
fraction of gradient elements over tol	81.4%	0.0%
conv1_conv/kernel grad max abs diff	0.5895	7.63e-06
worst per-tensor cosine similarity	-0.121	0.328

Per-tensor tolerance = 1e-5 + 1e-4 * max(|eager|, |graph|). With Grappler enabled the first conv layer's kernel gradient differs by 0.59 in max-abs terms and 161 of 214 captured backward tensors leave tolerance, with some tensors' gradients pointing in substantially different directions (cosine down to -0.12); disabling Grappler collapses the same comparison to 1/214 tensors and a 7.6e-6 conv1 gap.

What I am and am not claiming. The forward loss agreeing to 5.3e-7 (and bitwise with Grappler off) rules out an input or initial-weight mismatch. I am not claiming to know which arm is correct. I am reporting that a single backward step under tf.function(jit_compile=False) produces gradients that differ from eager on 161/214 tensors only when Grappler is enabled, which points at Grappler reordering the floating-point reductions in the backward graph.

Expected behavior

With identical inputs and identical initial weights, and a forward pass that agrees to 5.3e-7, I would expect the backward gradients to agree to within fp32 reassociation round-off as well — as they do when Grappler is disabled (1/214 tensors, conv1 gap 7.6e-6). If a 161/214-tensor backward divergence with Grappler enabled is expected reassociation from reduction reordering, I would appreciate confirmation; if not, it suggests a Grappler graph rewrite that changes backward numerics more than intended.

Standalone code to reproduce the issue

Save the following as a single file and run the command below it (CPU-only; CUDA_VISIBLE_DEVICES= is set inside the script). Each arm runs in its own child process and the driver prints one RESULT JSON line.

#!/usr/bin/env python3
"""Self-contained minimal reproducer: single backward step, eager vs
tf.function(jit_compile=False), per-tensor gradients, Grappler ON vs OFF.

Single CPU device, one deterministic synthetic batch, identical serialised
initial weights.  Each (arm, grappler) combination runs in its own child process
so the process-global Grappler meta-optimizer setting cannot leak between runs.

Reports, for Grappler ON and OFF separately: the per-tensor max_abs_diff of the
`conv1_conv/kernel` gradient, how many of the captured backward tensors exceed a
combined atol+rtol tolerance, the worst cosine similarity, and the fraction of
gradient elements over tolerance.

Usage:
  python tf_min_backward_grappler.py --model resnet50 --input-size 224 --batch 2
"""
import argparse, json, os, subprocess, sys, tempfile

os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # CPU-only

ATOL, RTOL = 1e-5, 1e-4


def build_model(model_name, num_classes, input_size):
    import tensorflow as tf
    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, name="resnet50_cls")


def run_arm(args, arm, grappler, weights_path, out_npz):
    import tensorflow as tf
    import numpy as np
    if grappler == "off":
        tf.config.optimizer.set_experimental_options({"disable_meta_optimizer": True})
    model = build_model(args.model, 10, args.input_size)
    model.load_weights(weights_path)
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    tv = model.trainable_variables
    x = tf.random.stateless_normal((args.batch, args.input_size, args.input_size, 3),
                                   seed=(args.seed, 0), dtype=tf.float32)
    y = tf.range(args.batch, dtype=tf.int32) % 10

    def fwd_bwd():
        with tf.GradientTape() as tape:
            logits = model(x, training=True)
            loss = loss_fn(y, logits)
        return tape.gradient(loss, tv), loss

    if arm == "graph":
        fwd_bwd = tf.function(fwd_bwd, jit_compile=False)
    grads, loss = fwd_bwd()
    arrs = {}
    for i, (v, g) in enumerate(zip(tv, grads)):
        path = str(getattr(v, "path", v.name)).replace("/", "_").replace(":", "_")
        arrs[f"{i:04d}_{path}"] = g.numpy()
    arrs["__loss__"] = np.array(float(loss))
    np.savez(out_npz, **arrs)


def save_init_weights(args, path):
    import tensorflow as tf
    tf.keras.utils.set_random_seed(args.seed)
    model = build_model(args.model, 10, args.input_size)
    _ = model(tf.zeros((1, args.input_size, args.input_size, 3)), training=False)
    model.save_weights(path)


def compare(npz_e, npz_g):
    import numpy as np
    e, g = np.load(npz_e), np.load(npz_g)
    keys = [k for k in e.files if k != "__loss__"]
    n_div = n_over = n_elem = 0
    worst_cos, conv1 = 1.0, None
    for k in keys:
        a, b = e[k].astype(np.float64).ravel(), g[k].astype(np.float64).ravel()
        diff = np.abs(a - b)
        tol = ATOL + RTOL * np.maximum(np.abs(a), np.abs(b))
        over = diff > tol
        n_over += int(over.sum()); n_elem += int(a.size)
        if over.any():
            n_div += 1
        na, nb = np.linalg.norm(a), np.linalg.norm(b)
        cos = float(a @ b / (na * nb)) if na > 0 and nb > 0 else 1.0
        worst_cos = min(worst_cos, cos)
        if "conv1_conv_kernel" in k:
            conv1 = float(diff.max())
    le = float(e["__loss__"]); lg = float(g["__loss__"])
    return {"forward_loss_eager": round(le, 8), "forward_loss_graph": round(lg, 8),
            "forward_loss_rel": round(abs(le - lg) / max(1e-12, (abs(le) + abs(lg)) / 2), 8),
            "tensors": len(keys), "divergent_tensors": n_div,
            "elem_over_tol_frac": round(n_over / max(1, n_elem), 4),
            "worst_cosine": round(worst_cos, 4),
            "conv1_conv_kernel_max_abs_diff": conv1}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="resnet50")
    p.add_argument("--input-size", type=int, default=224)
    p.add_argument("--batch", type=int, default=2)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--arm", default=None)
    p.add_argument("--grappler", default="on")
    p.add_argument("--weights", default=None)
    p.add_argument("--out", default=None)
    args = p.parse_args()

    if args.arm:
        run_arm(args, args.arm, args.grappler, args.weights, args.out)
        return

    with tempfile.TemporaryDirectory() as td:
        wpath = os.path.join(td, "init.weights.h5")
        save_init_weights(args, wpath)
        base = [sys.executable, os.path.abspath(__file__),
                "--model", args.model, "--input-size", str(args.input_size),
                "--batch", str(args.batch), "--seed", str(args.seed), "--weights", wpath]
        npzs = {}
        for grappler in ("on", "off"):
            for arm in ("eager", "graph"):
                o = os.path.join(td, f"g_{grappler}_{arm}.npz")
                subprocess.run(base + ["--arm", arm, "--grappler", grappler, "--out", o],
                               check=True, capture_output=True, text=True)
                npzs[(grappler, arm)] = o
        report = {
            "model": args.model, "input_size": args.input_size, "batch": args.batch,
            "grappler_on": compare(npzs[("on", "eager")], npzs[("on", "graph")]),
            "grappler_off": compare(npzs[("off", "eager")], npzs[("off", "graph")]),
        }
        print("RESULT " + json.dumps(report))


if __name__ == "__main__":
    main()

Command:

python tf_min_backward_grappler.py --model resnet50 --input-size 224 --batch 2 --seed 42

Relevant log output

RESULT {"model": "resnet50", "input_size": 224, "batch": 2, "grappler_on": {"forward_loss_eager": 1.57970905, "forward_loss_graph": 1.57970989, "forward_loss_rel": 5.3e-07, "tensors": 214, "divergent_tensors": 161, "elem_over_tol_frac": 0.8137, "worst_cosine": -0.1206, "conv1_conv_kernel_max_abs_diff": 0.5894842147827148}, "grappler_off": {"forward_loss_eager": 1.57970905, "forward_loss_graph": 1.57970905, "forward_loss_rel": 0.0, "tensors": 214, "divergent_tensors": 1, "elem_over_tol_frac": 0.0, "worst_cosine": 0.3276, "conv1_conv_kernel_max_abs_diff": 7.62939453125e-06}}