Spaces:
Sleeping
Sleeping
File size: 4,334 Bytes
15bcbe6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# Copyright 2022 Google.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Flax modules and functions for using external memory."""
from typing import Any, Optional, Tuple
from absl import logging
from flax import linen
import gin
import jax
from transformer import memory_layer
PRNGKey = Any
Shape = Tuple[int]
Dtype = Any
Array = Any
MemoryResource = Any
class MemoryManager:
"""Manages any external resources that may be required by external memory.
MemoryManager also functions as a factory, to create Flax modules that will
read and write to whatever external memory has been configured.
"""
def __init__(self,
batch_size: int,
mode: str,
num_heads: int,
key_size: int,
value_size: int,
database_size: Optional[int] = None,
dtype: Dtype = "float32",
off_device_memory: Optional[MemoryResource] = None):
"""Create a MemoryManager object.
A MemoryManager configures external memory, and is used as a factory to
construct flax modules that read or write to the memory.
Args:
batch_size: The number of separate documents in a batch.
mode: e.g. ("train", or "test")
num_heads: The number of transformer heads.
key_size: The length of the key vectors.
value_size: The length of the value vectors.
database_size: The total number of tokens in the database.
dtype: The datatype used for keys and values.
off_device_memory: An object which manages underlying SCAM memory.
If None, then the model will use on-device memory.
"""
self.batch_size = batch_size
self.mode = mode
self.num_heads = num_heads
self.key_size = key_size
self.value_size = value_size
self.database_size = database_size
self.dtype = dtype
self.off_device_memory = off_device_memory
def create_memory_layer(self) -> linen.Module:
"""Create a flax Module that implements external memory."""
num_datasets = (
self.batch_size * self.num_heads #
if self.off_device_memory is None #
else self.num_heads)
if self.off_device_memory is not None:
mem_layer = None
if mem_layer is None:
raise ValueError("Off-device memory is not supported at this time.")
return memory_layer.BatchedMemory(
mem_layer,
split_dimensions=(-2,),
)
else:
assert self.database_size is not None
mem_layer = memory_layer.MemoryOnTpu(num_datasets=num_datasets,
key_features=self.key_size,
value_features=self.value_size,
database_size=self.database_size,
dtype=self.dtype)
# Handle queries of shape [batch_size, seq_len, num_heads, kv_features]
return memory_layer.BatchedMemory(mem_layer,
split_dimensions=(0, -2))
@gin.configurable
def memory_on_tpu_factory(batch_size: int,
mode: str,
num_heads: int = gin.REQUIRED,
key_size: int = gin.REQUIRED,
value_size: int = gin.REQUIRED,
database_size: int = gin.REQUIRED,
dtype: Dtype = gin.REQUIRED) -> MemoryManager:
"""Implement SCAM memory on device."""
return MemoryManager(batch_size=batch_size,
mode=mode,
num_heads=num_heads,
key_size=key_size,
value_size=value_size,
database_size=database_size,
dtype=dtype,
off_device_memory=None)
|