Normalize device kind strings when looking up HBM memory.

This change removes spaces from device kind strings before looking up HBM memory values. This ensures that device kinds like "TPU 7x" and "TPU7x" are treated the same, preventing lookup failures due to inconsistent spacing.

PiperOrigin-RevId: 870223149

Author: samos123

CHANGELOG.md

@@ -22,6 +22,11 @@ handlers using `StepMetadata.item_handlers` and the global `HandlerTypeRegistry`
 if no args are provided.
 - `CompositeCheckpointHandler.metadata()` now returns `StepMetadata`.
 
+### Fixed
+
+- Fixed `get_device_memory` issue on TPU 7x devices where the device kind string
+was consistently reported without a space, causing a ValueError.
+
 ## [0.1.7] - 2022-03-29
 
 ### Added

checkpoint/orbax/checkpoint/_src/multihost/multislice.py

@@ -159,7 +159,10 @@ def get_device_memory() -> int:
       'NVIDIA B200': int(183e9),
       'NVIDIA B300 SXM6 AC': int(275e9),
   }
-  memory = hbm_memory.get(device.device_kind, None)
+  # Remove spaces from the device kind to make the lookup robust.
+  # For example, "TPU 7x" and "TPU7x" should both map to the same value.
+  normalized_hbm_memory = {k.replace(' ', ''): v for k, v in hbm_memory.items()}
+  memory = normalized_hbm_memory.get(device.device_kind.replace(' ', ''), None)
   if memory is None:
     raise ValueError(
         f'get_device_memory is not supported for {device.device_kind}.'