Just to make it explicit (you noticed, but for the next person looking here), clamping the quotient won’t work, but using
    with torch.no_grad():
        norm = w.norm(2, dim=0, keepdim=True).clamp(min=self._max_norm_val / 2)
        desired = torch.clamp(norm, max=self._max_norm_val)
        w *= (desired / norm)

should always work as expected and leave small norms alone.