Repository: pytorch
Issue: mutating a buffer to change its requires_grad-ness can cause RuntimeError
Labels: triaged,oncall: pt2,module: aotdispatch,module: pt2-dispatcher

repro (from @fmassa):
```
from itertools import chain
from typing import Any, Dict, Optional, Union
import torch
from torch import nn
from torch.nn import functional as F
import os
os.environ["FI_EFA_SET_CUDA_SYNC_MEMOPS"] = "0"
os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "10"
def _parse_slurm_node_list(s: str):
    import re
    nodes = []
    # Extract "hostname", "hostname[1-2,3,4-5]," substrings
    p = re.compile(r"(([^\[]+)(?:\[([^\]]+)\])?),?")
    for m in p.finditer(s):
        prefix, suffixes = s[m.start(2) : m.end(2)], s[m.start(3) : m.end(3)]
        for suffix in suffixes.split(","):
            span = suffix.split("-")
            if len(span) == 1:
                nodes.append(prefix + suffix)
            else:
                width = len(span[0])
                start, end = int(span[0]), int(span[1]) + 1
                nodes.extend([prefix + f"{i:0{width}}" for i in range(start, end)])
    return nodes
IS_DIST = True
if os.environ.get("RANK", ""):
    rank = int(os.environ.get("RANK"))
    local_rank = rank
    world_size = int(os.environ['WORLD_SIZE'])
elif 'SLURM_PROCID' in os.environ:
    rank = int(os.environ['SLURM_PROCID'])
    local_rank = int(os.environ["SLURM_LOCALID"])
    world_size = int(os.environ.get("SLURM_NTASKS"))
    os.environ["MASTER_ADDR"] = _parse_slurm_node_list(os.environ["SLURM_JOB_NODELIST"])[0]
    os.environ["MASTER_PORT"] = "24189"
else:
    IS_DIST = False
    local_rank = 0
if IS_DIST:
    pg_options = torch.distributed.ProcessGroupNCCL.Options(is_high_priority_stream=True)
    torch.distributed.init_process_group(backend="nccl", init_method="env://",
            world_size=world_size, rank=rank, pg_options=pg_options)
    torch.distributed.barrier()
    torch.cuda.set_device(local_rank)
def get_tensor_model_parallel_group():
    return torch.distributed.distributed_c10d._get_default_group()
def get_tensor_model_parallel_world_size():
    return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())
def get_tensor_model_parallel_group_name():
    return get_tensor_model_parallel_group().group_name
def reduce_gradient_across_mp_group(grad):
    torch.distributed.all_reduce(grad, group=get_tensor_model_parallel_group())
    return grad
def init_with_mp_rank_seed(weight, init_method):
    init_method(weight)
def all_to_all(x: torch.Tensor, split_dim: int, concat_dim: int):
    mp_size = get_tensor_model_parallel_world_size()
    if mp_size == 1:
        return x
    chunks = torch.cat(torch.tensor_split(x, mp_size, dim=split_dim), dim=0)

    output_split_sizes = [chunks.shape[0] // mp_size] * mp_size
    input_split_sizes = output_split_sizes
    group_name = get_tensor_model_parallel_group_name()
    output = torch.ops._c10d_functional_autograd.all_to_all_single(
        chunks,
        output_split_sizes,
        input_split_sizes,
        group_name,
    )
    output = torch.ops._c10d_functional.wait_tensor(output)
    output = torch.cat(torch.tensor_split(output, mp_size, dim=0), dim=concat_dim)
    return output
class AllReduceInBackward(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, group_name):
        ctx.group_name = group_name
        return x
    @staticmethod
    def backward(ctx, grad):
        grad = torch.ops._c10d_functional.all_reduce(grad, "sum", ctx.group_name)
        grad = torch.ops._c10d_functional.wait_tensor(grad)
        return grad, None, None
def all_reduce_in_backward(x):
    group_name = get_tensor_model_parallel_group_name()
    return AllReduceInBackward.apply(x, group_name)
class ExpertsChoiceMOE(torch.nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        ffn_dim_multiplier: float,
        multiple_of: int,
        non_linearity: str,
        init_depth: Optional[int],
        layer_id: int,
        # MOE specific parameters
        number_of_experts: int = 8,  # more experts = more sparse params = better quality (at expense of memory)
        capacity_factor: float = 1.0,  # capacity factor determines how many tokens each expert can choose
        auto_scale_F: bool = True,  # if true, rescales hidden_dim such that number of activates params is same as equivalent dense layer
        use_shared_expert: bool = True,  # if true, creats a determinstic shared expert to be activated alongside the routed experts
        pregate_moe: bool = False,  # if true, applies the gate multiplier to the input into routed experts instead of output. this can save memory on backward pass
        shared_expert_gate: bool = True,  # if true, learns a parameter that determines how to mix shared and routed outputs, but requires more memory in the backward pass
        eval_with_saved_stats: bool = False,  # if true, uses saved gate stats as thresholds for eval
        eval_threshold_std_mult: float = 0.0,  # adjust threshold by multiple of standard deviation (for elastic compute)
    ):
        super().__init__()
        self.layer_id = layer_id
        self.capacity_factor = capacity_factor
        self.use_shared_expert = use_shared_expert
        self.pregate_moe = pregate_moe
        self.shared_expert_gate = shared_expert_gate
        self.eval_with_saved_stats = eval_with_saved_stats
        self.eval_threshold_std_mult = eval_threshold_std_mult
        self.E = number_of_experts
        self.swiglu = non_linearity == "swiglu"
        # swiglu hidden dim factor multiplier (same #params as relu / gelu)
        if self.swiglu:
            hidden_dim = int(2 * hidden_dim / 3)
        # custom dim factor multiplier
        hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        if auto_scale_F:
            hidden_dim = int(hidden_dim / (capacity_factor + int(use_shared_expert)))
        # round hidden dimension to `multiple_of`
        hidden_dim += -hidden_dim % multiple_of
        mp_size = get_tensor_model_parallel_world_size()
        assert (
            number_of_experts % mp_size == 0
        ), f"number_of_experts ({number_of_experts}) must be divisible by mp_size ({mp_size})"
        e = number_of_experts // mp_size
        init_in_fn = lambda x: torch.nn.init.trunc_normal_(x)
        init_out_fn = lambda x: torch.nn.init.trunc_normal_(x)
        # moe layers
        self.moe_w_in_eDF = nn.Parameter(torch.empty(e, dim, hidden_dim))
        init_with_mp_rank_seed(self.moe_w_in_eDF, init_in_fn)
        self.moe_w_out_eFD = nn.Parameter(torch.empty(e, hidden_dim, dim))
        init_with_mp_rank_seed(self.moe_w_out_eFD, init_out_fn)
        self.router_DE = nn.Parameter(torch.empty(dim, number_of_experts))
        nn.init.normal_(self.router_DE, mean=0.0, std=0.8 * dim**-0.5)
        if use_shared_expert:
            self.w_in_shared_DF = nn.Parameter(torch.empty(dim, hidden_dim))
            init_in_fn(self.w_in_shared_DF)
            self.w_out_shared_FD = nn.Parameter(torch.empty(hidden_dim, dim))
            init_out_fn(self.w_out_shared_FD)
            if self.shared_expert_gate:
                self.shared_gate_1 = nn.Parameter(torch.full((1,), fill_value=0.0))
        if self.swiglu:
            self.moe_w_swiglu_eDF = nn.Parameter(torch.empty(e, dim, hidden_dim))
            init_with_mp_rank_seed(self.moe_w_swiglu_eDF, init_in_fn)
            if use_shared_expert:
                self.w_swiglu_DF = nn.Parameter(torch.empty(dim, hidden_dim))
                init_in_fn(self.w_swiglu_DF)
        # non-linearity
        self.non_linearity = {
            "relu": F.relu,
            "gelu": F.gelu,
            "swiglu": None,
            "srelu": lambda x: F.relu(x) ** 2,
            # "silu": F.silu,
            # "mish": F.mish,
            # "swish": swish,
        }[non_linearity]
        # sum, squared sum, count
        self.register_buffer("running_gate_stats_3E", torch.zeros(3, number_of_experts))
        self.running_gate_ema = 0.99
        # have to do it here before wrapping so we can actually access params
        self.repr_str = ""
        for n, p in chain(self.named_parameters(), self.named_buffers()):
            self.repr_str += f"{n}: {p.shape}\n"
    def _add_grad_reduce_hooks(self):
        if get_tensor_model_parallel_world_size() <= 1:
            return
        self.router_DE.register_hook(reduce_gradient_across_mp_group)
        if self.use_shared_expert:
            self.w_in_shared_DF.register_hook(reduce_gradient_across_mp_group)
            self.w_out_shared_FD.register_hook(reduce_gradient_across_mp_group)
            if self.shared_expert_gate:
                self.shared_gate_1.register_hook(reduce_gradient_across_mp_group)
            if self.swiglu:
                self.w_swiglu_DF.register_hook(reduce_gradient_across_mp_group)
    def forward(self, x_aD):
        a, D = x_aD.shape
        E = self.E
        # FIXME(fmassa): I replaced those backward hooks with all_reduce_in_backward. Could also have used DTensor
        # self._add_grad_reduce_hooks()
        # get router scores
        router_DE = all_reduce_in_backward(self.router_DE)
        router_scores_Ea = torch.einsum("aD,DE->Ea", x_aD, router_DE)
        router_scores_Ea = torch.sigmoid(router_scores_Ea)
        tokens_per_expert = int(a * self.capacity_factor / E)
        tokens_per_expert += -tokens_per_expert % 8  # round to multiple of 8
        stats_handle = None
        if self.training or not self.eval_with_saved_stats:
            router_scores_Eg, router_indices_Eg = torch.topk(
                router_scores_Ea, tokens_per_expert, dim=1
            )
            # update running stats
            self.running_gate_stats_3E.mul_(self.running_gate_ema)
            # FIXME(fmassa): I had to add the .detach() here otherwise the bwd would complain with aot_eager
            #min_thresh_E = router_scores_Eg.min(dim=1, keepdim=False).values.detach()
            min_thresh_E = router_scores_Eg.min(dim=1, keepdim=False).values
            self.running_gate_stats_3E[0] += min_thresh_E
            self.running_gate_stats_3E[1] += min_thresh_E**2
            self.running_gate_stats_3E[2] += 1
            if torch.distributed.is_initialized():
                stats_handle = torch.distributed.all_reduce(
                    self.running_gate_stats_3E,
                    op=torch.distributed.ReduceOp.AVG,
                    # async_op=True,
                    async_op=False,
                )
        else:
            count = self.running_gate_stats_3E[2]
            assert count > 0
            mean_E = self.running_gate_stats_3E[0] / count
            std_E = torch.sqrt(self.running_gate_stats_3E[1] / count - mean_E**2)
            threshold_E = mean_E + std_E * self.eval_threshold_std_mult
            router_scores_Eg = torch.where(
                router_scores_Ea >= threshold_E.unsqueeze(1),
                router_scores_Ea,
                torch.zeros_like(router_scores_Ea),
            )
            router_indices_Eg = (
                torch.arange(a, device=x_aD.device).view(1, -1).expand(E, -1)
            )
        routed_in_Eg_D = x_aD.gather(
            dim=0, index=router_indices_Eg.view(-1, 1).expand(-1, D)
        )
        if self.pregate_moe:
            routed_in_Eg_D = routed_in_Eg_D * router_scores_Eg.view(-1, 1)
        routed_in_eGD = all_to_all(
            routed_in_Eg_D.view(E, tokens_per_expert, D), split_dim=0, concat_dim=1
        )
        routed_hidden_eGF = torch.einsum(
            "eGD,eDF->eGF", routed_in_eGD, self.moe_w_in_eDF
        )
        if self.swiglu:
            swiglu_hidden_eGF = torch.einsum(
                "eGD,eDF->eGF", routed_in_eGD, self.moe_w_swiglu_eDF
            )
            routed_middle_eGF = F.silu(routed_hidden_eGF) * swiglu_hidden_eGF
        else:
            routed_middle_eGF = self.non_linearity(routed_hidden_eGF)
        routed_out_eGD = torch.einsum(
            "eGF,eFD->eGD", routed_middle_eGF, self.moe_w_out_eFD
        )
        routed_out_EgD = all_to_all(routed_out_eGD, split_dim=1, concat_dim=0)
        if not self.pregate_moe:
            routed_out_EgD = routed_out_EgD * router_scores_Eg.view(
                E, tokens_per_expert, 1
            )
        routed_out_aD = torch.zeros_like(x_aD)
        router_indices_Eg_D = router_indices_Eg.view(-1, 1).expand(-1, D)
        routed_out_aD.scatter_add_(
            dim=0, index=router_indices_Eg_D, src=routed_out_EgD.view(-1, D)
        )
        out_aD = routed_out_aD
        # try to overlap this computation with comms from above
        if self.use_shared_expert:
            w_in_shared_DF = all_reduce_in_backward(self.w_in_shared_DF)
            shared_middle_aF = x_aD @ w_in_shared_DF
            if self.swiglu:
                w_swiglu_DF = all_reduce_in_backward(self.w_swiglu_DF)
                swiglu_hidden_aF = x_aD @ w_swiglu_DF
                shared_middle_aF = F.silu(shared_middle_aF) * swiglu_hidden_aF
            else:
                shared_middle_aF = self.non_linearity(shared_middle_aF)
            w_out_shared_FD = all_reduce_in_backward(self.w_out_shared_FD)
            shared_out_aD = shared_middle_aF @ w_out_shared_FD
            if self.shared_expert_gate:
                shared_gate_1 = all_reduce_in_backward(self.shared_gate_1)
                shared_gate_1 = torch.sigmoid(shared_gate_1)
                out_aD = shared_out_aD * shared_gate_1 + routed_out_aD * (
                    1 - shared_gate_1
                )
            else:
                out_aD = routed_out_aD + shared_out_aD
        # FIXME(fmassa): I made the all_reduce above be sync, so this is not needed anymore
        # if stats_handle:
        #     stats_handle.wait()
        return out_aD
    def extra_repr(self) -> str:
        return self.repr_str
dim = 512
ffn_exp = 4.0
ffn_dim_multiplier = 1.0
multiple_of = 256
model = ExpertsChoiceMOE(
    dim=dim, hidden_dim=int(ffn_exp * dim),
    ffn_dim_multiplier=ffn_dim_multiplier,
    multiple_of=multiple_of, non_linearity="swiglu",
    init_depth=None, layer_id=0)
dtype = torch.bfloat16
device = torch.device("cuda")
model.to(device, dtype)
#print(model)
x = torch.rand(4096, 512, dtype=dtype, device=device)
#out = model(x)
#out.sum().backward()
print("pass 0")
#model = torch.compile(model, backend="aot_eager")
model = torch.compile(model)
out = model(x)
out.sum().backward()
out = model(x)
out.sum().backward()
print("pass")
```

Note that the error only happens when invoking the fw/bw twice. It fails with:
```
[rank0]:   File "/home/hirsheybar/local/b/pytorch/torch/autograd/graph.py", line 825, in _engine_run_backward
[rank0]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:   File "/home/hirsheybar/local/b/pytorch/torch/autograd/function.py", line 307, in apply
[rank0]:     return user_fn(self, *args)
[rank0]:   File "/home/hirsheybar/local/b/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1692, in backward
[rank0]:     all_args = CompiledFunction._backward_prologue(ctx, *flat_args)
[rank0]:   File "/home/hirsheybar/local/b/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1838, in _backward_prologue
[rank0]:     ctx_saved_tensors = ctx.saved_tensors
[rank0]: RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
```

In the repro code, we have a buffer (which does not require grad), that we mutate with a tensor that requires grad. Tweaking the repro to use:
```
min_thresh_E = router_scores_Eg.min(dim=1, keepdim=False).values.detach()
```
fixes the error under compile, although this change wasn't necessary under eager

cc @ezyang @chauhang @penguinwu @zou3519 @yf225