I want to confirm there will be total 48g memory while I use nvlink to connect two 3090.

That sounds right, but since this is a GPU hardware spec-related question, I would ask the NVIDIA team directly.

If rtx3090 supports this feature, how should I change my pytorch code?

With 2 GPUs and nvlink connecting them, I would use DistributedDataParallel (DDP) for training. At a high level, you can spawn 2 CPU processes, 1 for each GPU, and create a NCCL Process Group to have fast data transfer between the 2 GPUs. Then you can simply wrap your model with DDP and train.
Here are some references in terms of how to implement:


docs: https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel 56


Implementation: https://pytorch.org/docs/master/notes/ddp.html#example 24


Example: https://github.com/pytorch/examples/tree/master/distributed/ddp 34