Merge pull request #2597 from keshavb96/b300_mem_key

mridul-sahu · web-flow · commit 79b3896e011a · 2025-11-18T22:31:31.000+05:30
add memory key for B300
diff --git a/checkpoint/orbax/checkpoint/_src/multihost/multislice.py b/checkpoint/orbax/checkpoint/_src/multihost/multislice.py
@@ -156,6 +156,7 @@ def get_device_memory() -> int:
       'NVIDIA H100 80GB HBM3': int(80e9),
       'NVIDIA H200': int(144e9),
       'NVIDIA B200': int(183e9),
+      'NVIDIA B300 SXM6 AC': int(275e9)
   }
   memory = hbm_memory.get(device.device_kind, None)
   if memory is None:

Original file line number	Diff line number	Diff line change
`@@ -156,6 +156,7 @@ def get_device_memory() -> int:`
`156`	`156`	`'NVIDIA H100 80GB HBM3': int(80e9),`
`157`	`157`	`'NVIDIA H200': int(144e9),`
`158`	`158`	`'NVIDIA B200': int(183e9),`
	`159`	`+ 'NVIDIA B300 SXM6 AC': int(275e9)`
`159`	`160`	`}`
`160`	`161`	`memory = hbm_memory.get(device.device_kind, None)`
`161`	`162`	`if memory is None:`