nccl
This commit is contained in:
parent
96919dc0b5
commit
b198b6fdca
|
|
@ -0,0 +1,17 @@
|
||||||
|
import os
|
||||||
|
import torch
|
||||||
|
import torch.distributed as dist
|
||||||
|
|
||||||
|
def main():
|
||||||
|
dist.init_process_group("nccl")
|
||||||
|
local_rank = int(os.environ["LOCAL_RANK"])
|
||||||
|
torch.cuda.set_device(local_rank)
|
||||||
|
|
||||||
|
print(f"Rank {dist.get_rank()} initialized on GPU {local_rank}")
|
||||||
|
|
||||||
|
x = torch.ones(10).cuda()
|
||||||
|
dist.all_reduce(x)
|
||||||
|
print(f"Rank {dist.get_rank()} result: {x[0].item()}")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
|
@ -365,8 +365,8 @@ spec:
|
||||||
template:
|
template:
|
||||||
metadata:
|
metadata:
|
||||||
labels: { app: nccl }
|
labels: { app: nccl }
|
||||||
annotations:
|
# annotations:
|
||||||
k8s.v1.cni.cncf.io/networks: hostdevice-net # HostDeviceNetwork (RDMA)
|
# k8s.v1.cni.cncf.io/networks: hostdevice-net # HostDeviceNetwork (RDMA)
|
||||||
spec:
|
spec:
|
||||||
# hostNetwork: true
|
# hostNetwork: true
|
||||||
# dnsPolicy: ClusterFirstWithHostNet
|
# dnsPolicy: ClusterFirstWithHostNet
|
||||||
|
|
@ -392,15 +392,17 @@ spec:
|
||||||
resources:
|
resources:
|
||||||
limits:
|
limits:
|
||||||
nvidia.com/gpu: "8"
|
nvidia.com/gpu: "8"
|
||||||
nvidia.com/hostdev: "8" # SR-IOV Device Plugin이 광고한 RDMA NIC 리소스
|
# nvidia.com/hostdev: "8" # SR-IOV Device Plugin이 광고한 RDMA NIC 리소스
|
||||||
requests:
|
requests:
|
||||||
nvidia.com/gpu: "8"
|
nvidia.com/gpu: "8"
|
||||||
nvidia.com/hostdev: "8"
|
# nvidia.com/hostdev: "8"
|
||||||
volumeMounts:
|
volumeMounts:
|
||||||
- name: scripts
|
- name: scripts
|
||||||
mountPath: /config_scripts
|
mountPath: /config_scripts
|
||||||
- name: shared-memory
|
- name: shared-memory
|
||||||
mountPath: /dev/shm
|
mountPath: /dev/shm
|
||||||
|
- name: infiniband
|
||||||
|
mountPath: /dev/infiniband
|
||||||
volumes:
|
volumes:
|
||||||
- name: scripts
|
- name: scripts
|
||||||
configMap:
|
configMap:
|
||||||
|
|
@ -410,5 +412,9 @@ spec:
|
||||||
emptyDir:
|
emptyDir:
|
||||||
medium: Memory
|
medium: Memory
|
||||||
sizeLimit: 118541097369600m
|
sizeLimit: 118541097369600m
|
||||||
|
- name: infiniband
|
||||||
|
hostPath:
|
||||||
|
path: /dev/infiniband
|
||||||
|
type: Directory
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue