| 496 | return array_serialization.get_tensorstore_spec(ckpt_path) |
| 497 | |
| 498 | def _get_spec(self, step: int, state: NestedTensor, ckpt_dir: str) -> CheckpointSpec: |
| 499 | spec = self.CheckpointSpec( |
| 500 | index=[], |
| 501 | storage_paths=[], |
| 502 | tensorstore_specs=[], |
| 503 | shapes=[], |
| 504 | dtypes=[], |
| 505 | shardings=[], |
| 506 | gda_values=[], |
| 507 | tf_ckpt_map={}, |
| 508 | python_ckpt_map={}, |
| 509 | ) |
| 510 | |
| 511 | mesh = thread_resources.env.physical_mesh |
| 512 | if not mesh.shape: |
| 513 | raise RuntimeError( |
| 514 | "Checkpoint restoration must take place within the context of a Mesh" |
| 515 | ) |
| 516 | spec.index = [("step", step)] |
| 517 | for path, value in utils.flatten_items(state, separator="/"): |
| 518 | if isinstance(value, (Tensor, TensorSpec)): |
| 519 | logging.vlog( |
| 520 | 3, "Adding array value %s %s(%s)", type(value), value.dtype, value.shape |
| 521 | ) |
| 522 | dtype = getattr(value.dtype, "dtype", value.dtype) |
| 523 | spec.index.append((path, {"dtype": str(dtype), "shape": str(tuple(value.shape))})) |
| 524 | gda_path = os.path.join(ckpt_dir, "gda", path) |
| 525 | spec.storage_paths.append(gda_path) |
| 526 | spec.tensorstore_specs.append(self._spec_from_path(gda_path)) |
| 527 | spec.shapes.append(value.shape) |
| 528 | spec.dtypes.append(dtype) |
| 529 | if isinstance(value, Tensor): |
| 530 | spec.gda_values.append(value) |
| 531 | spec.shardings.append(value.sharding) |
| 532 | else: |
| 533 | spec.shardings.append( |
| 534 | jax.sharding.NamedSharding( |
| 535 | mesh, |
| 536 | ( |
| 537 | jax.sharding.PartitionSpec() |
| 538 | if value.mesh_axes is None |
| 539 | else value.mesh_axes |
| 540 | ), |
| 541 | ) |
| 542 | ) |
| 543 | elif isinstance(value, (tf.data.Iterator, ElasticDatasetIterator)): |
| 544 | logging.vlog(3, "Adding value (%s) to tf_ckpt_map", value) |
| 545 | spec.index.append((path, str(type(value)))) |
| 546 | spec.tf_ckpt_map[path] = value |
| 547 | elif isinstance(value, PythonSavable): |
| 548 | spec.index.append((path, str(type(value)))) |
| 549 | spec.python_ckpt_map[path] = value |
| 550 | else: |
| 551 | logging.vlog(3, "Adding value (%s) to index", value) |
| 552 | spec.index.append((path, value)) |
| 553 | return spec |
| 554 | |
| 555 | def save_to_dir( |