@Karl, your question of repeated down-scaling of the weights (without restoring them after each forward) is well-placed. I haven't seen this before and doubt it could make the weights vanish (exponential down-scaling).
However, scaling weights and not activations/outputs could be a valid choice to avoid loss of precisions in inference in the following case (especially if training is done with BFloat16):

• Assume you want to scale the activation/output of a certain layer by a certain scalar (not necessarily a dynamic one as in L2 normalization) during training. Assume training is finished, model is saved and now you want to load it and infer with it using classical pytorch nn.Module forward call.
  Because this method doesn't have scaling, you would need to fuse the scalar with the weights after training and before saving the model (weights*=scalar). This will introduce rounding errors because in training you were forwarding like this: (W.T * x)*scalar, whereas after, in inference, you'll be forwarding like this: (W.T * scalar) * x. Where x is the input of the scaled layer.
  

This loss of precision can be significant if weights have low values (~1e-3, 1e-4 for ex). In which case, imprecisions on some values could reach 5-10%.
I tried scaling weights in a certain context (llm pretraining) then doing the forward pass, then restoring them by allocating to self.weight.data the previously cloned (unscaled) weight data. Because if I allocate the new scaled-weights to self.weight it'll break as it needs a torch.nn.Parameter.parameter and not a torch.Tensor. And I can't scale weights in-place because in-place modification of the graph's leaf varibales is impossible.
However, this breaks the autograd graph connections in a way I couldn't understand. Next layer would loose it's grad attribute somehow.

It would be very nice if someone has an idea on how to scale the weights before forwarding and backwarding by the layer then restoring the non-scaled weights again for the next forward/backward, without breaking the autograd graph neither making the weights become torch.Tensor and this loose their parameter class.