get_rib_acts is batch-size dependent (floating point error)

[stefan-apollo]

get_rib_acts_test runs with atol = 0 on Pythia, according to the comments in the test_pythia_14m_build_graph test. However, this seems to depend on the batch size. In test_pythia_14m_build_graph_jacobian on PR #294 I found it runs with atol = 0 when batch_size=900 but not when batch_size=1800.

Not sure if this is boring and expected, or worth investigating. Probably the former.

def test_pythia_14m_build_graph_jacobian():
    atol = 0  # Works with 0 for batch_size 900 but not 1800
    updates = [
        # Runs in around 75s on a5000
        {"basis_formula": "jacobian"},
        {"dataset": {"return_set_n_samples": 1}},
        {"dataset": {"n_ctx": 2}},
        {"batch_size": 1800},
        {"node_layers": ["ln2.1", "mlp_out.5", "unembed"]},
        {"calculate_edges": True},
        {"edge_formula": "stochastic"},
        {"n_stochastic_sources": 1},
    ]
    config = get_pythia_config(*updates)
    results = graph_build_test(config=config, atol=atol)
    get_rib_acts_test(results, atol=0)


test_pythia_14m_build_graph_jacobian()

2024-01-22 15:18:30 - INFO - Time to calculate edges: 0.17 minutes
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Loaded pretrained model pythia-14m into HookedTransformer
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Batches through entire model: 100%|████████████| 99/99 [00:00<00:00, 116.50it/s]
Traceback (most recent call last):
  File "/mnt/ssd-interp/stefan/rib/tests/test_build_graph.py", line 231, in <module>
    test_pythia_14m_build_graph_jacobian()
  File "/mnt/ssd-interp/stefan/rib/tests/test_build_graph.py", line 228, in test_pythia_14m_build_graph_jacobian
    get_rib_acts_test(results, atol=0)
  File "/mnt/ssd-interp/stefan/rib/tests/test_build_graph.py", line 154, in get_rib_acts_test
    assert_is_close(utils_rib_acts, test_rib_acts, atol=atol, rtol=1e-5)
  File "/mnt/ssd-interp/stefan/rib/tests/utils.py", line 140, in assert_is_close
    assert_close(
  File "/mnt/ssd-interp/stefan/rib-env/lib/python3.10/site-packages/torch/testing/_comparison.py", line 1520, in assert_close
    raise error_metas[0].to_error(msg)
AssertionError: Tensor-likes are not close!

Mismatched elements: 2 / 804864 (0.0%)
Greatest absolute difference: 1.1368683772161603e-13 at index (358, 1, 255) (up to 0 allowed)
Greatest relative difference: 9.613073780341265e-05 at index (358, 1, 255) (up to 1e-05 allowed)