@torch.compile materializes broadcasted tensors, resulting in huge allocations on GPU

[cbuchner1]

Hi,

after running a memory trace on the End 2 End example notebook with the techniques here, I am seeing a huge regression in peak GPU memory usage between Sionna 1.2.2 and 2.0.1

The scenario was changed to num_ofdm_sym = 14, otherwise it is verbatim (but run in a standalone script). GPU memory usage is determined with torch.cuda.max_memory_allocated() and tf.config.experimental.get_memory_info('GPU:0')['peak'] respectively.

Num Slots	Sionna 1.2.2 Runtime	Sionna 1.2.2 GPU Peak Memory	Sionna 2.0.1 Runtime	Sionna 2.0.1 GPU Peak Memory
1	00:02:19	1.3757 GiB	00:01:27	18.3141 GiB
1000	00:02:35	2.0903 GiB	00:02:33	19.0327 GiB
10000	00:09:23	2.1685 GiB	00:16:50	19.0961 GiB

I am finding that the pyTorch graph compiler is a lot less likely to eliminate broadcast tensors from being materialized by fusing expressions than TensorFlow. In other words, the size of simulation scenarios in the SLS is suffering from restrictions that I haven't observed with TensorFlow and Sionna 1.x

Running that End2End downlink example, I am seeing two places in Sionna code that results in huge allocation spikes, using torch 2.10.0

1. src/sionna/phy/channel/utils.py line 297
   h_f = h * e

This broadcast h * e expands to P times the size of the reduced product h_f = h_f.sum(dim=-3) by being materialized on the GPU with P representing the number of paths used by the channel model (e.g. often 24 for UMa).

2. src/sionna/phy/ofdm/precoding.py line 352
   h_eff = h @ g

For num_ofdm_symbols = 14 in that End 2 End example, the input and output tensors have these shapes

  h shape:     (1, 210, 21, 14, 128,  1,  12) ~0.707 GiB
  g shape:     (1,   1, 21, 14, 128, 12,  10) ~0.034 GiB
  h_eff shape: (1, 210, 21, 14,  128, 1,  10) ~0.589 GiB
  h_eff shape: (1, 210,  1, 21,  10, 14, 128) ~0.589 GiB  (after permute)

The broadcast performed during the multiplication is fully materialized on the GPU even in code running under @torch.compile and requires 210 times the size of the input tensor g (because it's being broadcast 210 times), making a ~7GiB allocation.

Would the Sionna team consider deploying workarounds to prevent the materialization of such broadcasts, or do you consider this a problem of pyTorch which should be fixed on their end?

[cbuchner1]

Possible workarounds to reduce the memory spikes are

1. This sequence of statements in cir_to_ofdm_channel()
   h = a.unsqueeze(-1)
   h_f = h * e
   h_f = h_f.sum(dim=-3)
   can be replaced by a single statement
   h_f = torch.einsum("...pn,...puf->...nf", a, e)

The above modification is a suggestion by gpt-5.4-high, and it claims that "I verified the new version is lint-clean, matches the old implementation numerically for both tau shapes supported by cir_to_ofdm_channel(), and still works under torch.compile.".

2. In compute_effective_channel() of PrecodedChannel as well as _compute_effective_channel() of RZFPrecoder
   the sequence of statements
   g = g.unsqueeze(1)
   h_eff = h @ g
   can be replaced by a single statement
   h_eff = torch.einsum("brtofxy,btofys->brtofxs", h, g)
   completely avoiding materialization of the broadcasted version of g

It does greatly reduce the output of torch.cuda.max_memory_allocated() for me, but I hate how much less readable this version is.

I do get a minor speedup and a nearly 4x reduction in peak GPU memory usage with the E2E sample, for 14 OFDM symbols and 100 slots simulation duration. Results are very similar (different rounding, but no major changes)

Before patch:
torch.cuda.max_memory_allocated(): 18.3148 GiB user 1m32.461s
After patch:
torch.cuda.max_memory_allocated(): 4.9795 GiB user 1m27.398s

Note that TensorFlow also has a tf.einsum() function and when running Sionna v1.x code in Eager mode this also materializes such large broadcast tensors. It may be worthwhile to also apply this workaround to Sionna v1.x.

[cbuchner1]

I am finding another spot that creates an allocation peak

_step_11_nlos() allocates a ~2.15 GiB h_full temporary in my scenario, then another ~2.15 GiB scaled copy before the first one is freed.
2.15 GiB at channel_coefficients.py:866
another 2.15 GiB at channel_coefficients.py:875

A second large hotspot is _step_11_reduce_nlos(), which duplicates h_full again via gather (channel_coefficients.py:931):

Here's a chart showing the two peaks from lines 875 and 931

My proposed solution to fix the first peak is computing h_full like this (no need to compute power_scaling explicitly)

        # b=batch, x=tx, y=rx, c=cluster, r=ray, m=rx antenna, n=tx antenna, t=time
        h_full = torch.einsum(
            "bxycrmn,bxycrmn,bxyc,bxycrt->bxycrmnt",
            h_field,
            h_array,
            torch.sqrt(rays.powers / h_doppler.shape[4]),
            h_doppler,
        )

I was able to eradicate the 2nd memory peak by first computing h_sub_cl_1,2,3 and then immediately contracting h_full along dim=4 before computing h_weak via gather.

            h_strong = torch.gather(h_full, dim=3, index=idx)

            # Sum specific rays for each sub-cluster
            h_sub_cl_1 = h_strong[:, :, :, :, self._sub_cl_1_ind, ...].sum(dim=4)
            h_sub_cl_2 = h_strong[:, :, :, :, self._sub_cl_2_ind, ...].sum(dim=4)
            h_sub_cl_3 = h_strong[:, :, :, :, self._sub_cl_3_ind, ...].sum(dim=4)

            # The other clusters are the weak clusters
            h_full = h_full.sum(dim=4) # early contraction saves memory
            strongest_rest = strongest_clusters[..., 2:]
            idx = strongest_rest.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
            idx = idx.expand(-1, -1, -1, -1, h_full.shape[4], h_full.shape[5], h_full.shape[6])
            h_weak = torch.gather(h_full, dim=3, index=idx)

torch.cuda.max_memory_allocated(): 3.8380 GiB

I wonder if a GPU memory peak reduction could be a viable task given the AI telco engineer.
https://github.com/NVlabs/the-ai-telco-engineer
It could optimize the peak GPU memory utilization downwards with the secondary objective of maintaining numerical results and execution speed - while ingesting a saved memory trace (a pickle file) to find code locations to inspect further.

[cbuchner1]

The get_pathloss() function in sionna/sys/utils.py does abs() followed by raising to the power of 2.

rx_power = torch.abs(h_freq).pow(2).to(dtype)
rx_power = rx_power.mean(dim=(-1, -3, -5))

This is inefficient, as it generates a huge temporary tensor rx_power and involves sqrt() and pow() functions and the conversion to dtype requires both the real valued and the complex valued tensor to briefly co-exist in memory.

Based on inspecting GPU memory traces, I am finding that computing power values like the suggestion below creates views into h_freq that do not explicitly materialize as temporaries and avoid GPU memory usage. Furthermore expensive sqrt() and pow() functions are avoided, so it may also be faster.

rx_power = (h_freq.real.square() + h_freq.imag.square()).mean(dim=(-1, -3, -5)).to(dtype)

The sub-optimal pattern .abs().pow(2) is also found elsewhere in Sionna code, e.g.

channel_gain = torch.abs(h).pow(2)
and
safe_pilots_expanded.abs().pow(2)

and it's unclear if the following frequently occuring pattern also has efficiency problems (other math libraries tend to offer an abs2() function to return the squared magnitude of matrix/tensor elements)

.abs().square()

[cbuchner1]

Computing rx covariance matrices is done like h @ h.mT is done in two places in src/sionna/phy/ofdm/equalization.py.
The problem is it takes a buffer twice the size of h for doing its work (reasons unclear). As h is the effective channel tensor, this uses enormous temporary storage.

The subroutine below does the same job, but computes the products and contraction using scalar views of the complex channel tensor. This avoids creating explicit copies except for the partial contraction results. I also tried doing it in complex domain like rxcov = (h.unsqueeze(-2) * h.unsqueeze(-3).conj()).sum(dim=-1) but this again would require significant work buffers. So I am working around inefficiencies of some complex pytorch operations by expressing them as real operations.

def memory_efficient_rx_covariance(
    h: torch.Tensor
) -> torch.Tensor:
    r"""Compute ``h @ h.mH`` using minimal temporary workspace.

    :param h: [..., num_rows, num_cols], `torch.complex`.
        Input tensor for the covariance computation.

    :output: [..., num_rows, num_rows], `torch.complex`.
        Covariance matrix computed as ``h @ h.mH``.
    """
    if h.dim() < 2:
        raise ValueError("`h` must have rank >= 2.")

    rxcov = torch.empty((*h.shape[:-1], h.shape[-2]), dtype=h.dtype, device=h.device)
    rxcov.real.copy_((h.real.unsqueeze(-2) * h.real.unsqueeze(-3)).sum(dim=-1))
    rxcov.real.add_((h.imag.unsqueeze(-2) * h.imag.unsqueeze(-3)).sum(dim=-1))
    rxcov.imag.copy_((h.imag.unsqueeze(-2) * h.real.unsqueeze(-3)).sum(dim=-1))
    rxcov.imag.sub_((h.real.unsqueeze(-2) * h.imag.unsqueeze(-3)).sum(dim=-1))
    return rxcov

In equalization.py it could be called like

        # was s_inf = h_dt_undesired @ h_dt_undesired.mT
        s_inf = memory_efficient_rx_covariance(h_dt_undesired)

or

        # was s_inf = h_eff_undesired @ h_eff_undesired.mT
        s_inf = memory_efficient_rx_covariance(h_eff_undesired)