※2020/10/24日時点の情報です。
10/28日にPyTorch 1.7が公開され，それに伴いCUDA11に対応・RTX3090でも利用可能となっています。
そのため，単にPyTorch 1.7に上げればRTX3080/3090でも使えます。
apexのインストールは，以下で出来ました。
https://github.com/NVIDIA/apex/issues/988#issuecomment-726343453
以下は古い情報です。
課題
PyTorch 1.6では，現状RTX 3080/3090に対応しておらず，以下のようなエラーを吐きます。
GeForce RTX 3090 with CUDA capability sm_86 is not compatible with the current PyTorch installation.
解決方法
対処方法は，開発版のPyTorch 1.8を入れることです。
pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cu110/torch_nightly.html
cu110となっていますが，cuda11.1を使っていても動くので問題ありません。
バージョン情報
NVIDIA-SMI 455.23.05 Driver Version: 455.23.05 CUDA Version: 11.1/Ubuntu 20.04
結果
以前は2070 SUPERを使っていたのですが，学習速度が3倍程度早くなりました。あと，jaxは普通に動きました。