def extend_step_fn(states, ids): with base_layer.JaxContext.new_context( prng_key=base_layer.next_prng_key(), global_step=global_step) as jax_context: jax_context.bind(self.model, self.model.vars_to_flax_vars(model_theta), [base_layer.SCOPE_AUX_LOSS]) new_states, xent = self.model.extend_step(states, ids) return new_states, xent.logits
def _dropout(self, inputs: JTensor, noise_shape: List[int]) -> JTensor: p = self.params if noise_shape is None: noise_shape = inputs.shape prng_key = base_layer.next_prng_key() keep_prob = p.keep_prob assert keep_prob > 0.0 random_nums = keep_prob + jax.random.uniform( prng_key, noise_shape, inputs.dtype, minval=0.0, maxval=1.0) binary_mask = jnp.floor(random_nums) return inputs * binary_mask / keep_prob
def decode(self, input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]: """Greedy decodes the input_batch. Args: input_batch: The input batch, with fields like `.ids`. Returns: - metrics, a NestedMap containing str keys and (metrics, weight) pairs. - A NestedMap like `input_batch`, with `.prefix_lengths` (vector of randomly generated ints indicating the lengths of prefixes for each row), and `.output_ids` (matrix of int ids with the decoded output). """ p = self.params if p.decoder.seqlen <= 0: raise ValueError('Must set p.decoder.seqlen > 0, current value = ' f'{p.decoder.seqlen}') batch_size = input_batch.ids.shape[0] maxval = jnp.sum(1 - input_batch.paddings, axis=1).astype(jnp.int32) minval = jnp.minimum(maxval, p.decoder.min_prefix_len) prefix_lengths = jax.random.randint(base_layer.next_prng_key(), [batch_size], minval, maxval + 1, input_batch.ids.dtype) decoder_state = self.lm.init_states( target_batch_size=batch_size, target_max_length=p.decoder.seqlen) global_step = base_layer.cur_global_step() lm_theta = self.lm.local_theta() def extend_step_fn(states, ids): with base_layer.JaxContext.new_context( prng_key=base_layer.next_prng_key(), global_step=global_step) as jax_context: jax_context.bind(self.lm, self.lm.vars_to_flax_vars(lm_theta), [base_layer.SCOPE_AUX_LOSS]) new_states, xent = self.lm.extend_step(states, ids) return new_states, xent.logits result = greedy_decode( extend_step_fn, decoder_state, input_batch.ids, input_batch.paddings, p.decoder.seqlen, max_decode_steps=p.decoder.max_decode_steps, prefix_lengths=prefix_lengths, eos_id=p.decoder.eos_id) result.update(input_batch) metrics = NestedMap( num_decoded=(jnp.array(0.0, jnp.float32), jnp.array(batch_size, jnp.float32))) return metrics, result
def fprop(self, inputs: JTensor, paddings: JTensor) -> Tuple[JTensor, JTensor]: """Applies data augmentation by randomly masking values in the spectrum. Args: inputs: A tensor of shape [batch, length, channels]. paddings: A 0/1 tensor of shape [batch, length]. Returns: A pair <new_inputs, mask>: new_inputs: A tensor of shape [batch, length, channels]. paddings: A 0/1 tensor of shape [batch, length]. """ lengths = jnp.einsum('bh->b', 1 - paddings).astype(jnp.int32) prng_key = base_layer.next_prng_key() inputs = self._time_mask(inputs, lengths, global_seed=prng_key) prng_key = base_layer.next_prng_key() inputs = self._frequency_mask(inputs, global_seed=prng_key) return inputs, paddings
def _apply_zoneout(self, state0: NestedMap, inputs: NestedMap, new_c: JTensor, new_m: JTensor) -> NestedMap: """Apply Zoneout and returns the updated states.""" p = self.params if p.zo_prob > 0.0: c_random_uniform = jax.random.uniform(base_layer.next_prng_key(), new_c.shape) m_random_uniform = jax.random.uniform(base_layer.next_prng_key(), new_m.shape) else: c_random_uniform = None m_random_uniform = None new_c = self._zoneout_internal(state0.c, new_c, inputs.padding, p.zo_prob, self.do_eval, c_random_uniform) new_m = self._zoneout_internal(state0.m, new_m, inputs.padding, p.zo_prob, self.do_eval, m_random_uniform) return NestedMap(m=new_m, c=new_c)
def body_fprop(self, per_stage_inputs: JTensor, *per_stage_args, **per_stage_kwargs) -> NestedJTensor: """Runs the fprop function of the stages.""" p = self.params if p.mesh_axis_names is not None: def annotate(x): unconstrained_dims = list(range(1, x.ndim)) dims_mapping = (p.weight_split_dims_mapping.stages + [None] * (x.ndim - 1)) return base_layer.maybe_shard(x, dims_mapping, p.mesh_axis_names, unconstrained_dims) per_stage_inputs = jax.tree_map(annotate, per_stage_inputs) per_stage_args = jax.tree_map(annotate, per_stage_args) per_stage_kwargs = jax.tree_map(annotate, per_stage_kwargs) prng_key = base_layer.next_prng_key() global_step = base_layer.cur_global_step() # vmap self.body.fprop to get a leading stage dimension to handle per_stage # inputs and args. def _wrapped_fn(theta, per_stage_inputs, *per_stage_args, **per_stage_kwargs): with base_layer.JaxContext.new_context( prng_key=prng_key, global_step=global_step) as jax_ctx: jax_ctx.bind(self.body, self.body.vars_to_flax_vars(theta), [base_layer.SCOPE_AUX_LOSS]) res = self.body.fprop(per_stage_inputs, *per_stage_args, **per_stage_kwargs) summaries = base_layer.all_summaries() return res, summaries res, summaries = jax.vmap(_wrapped_fn)(self.body.local_theta(), per_stage_inputs, *per_stage_args, **per_stage_kwargs) self._forward_summary(summaries) return res
def _drop_connect(self, inputs: JTensor) -> JTensor: """Drops the entire residual layer with given survival probability. Args: inputs: input `.JTensor` which is on the residual branch which is dropped. Returns: Dropped out inputs. """ if self.do_eval: return inputs # Compute tensor. prng_key = base_layer.next_prng_key() batch_size = inputs.shape[0] shape = [batch_size] + [1] * (len(inputs.shape) - 1) random_tensor = self.params.survival_prob + jax.random.uniform( prng_key, shape, dtype=inputs.dtype) binary_tensor = jnp.floor(random_tensor) # Unlike conventional way that multiply survival_prob at test time, here we # divide survival_prob at training time, such that no additional compute is # needed at test time. output = inputs / self.params.survival_prob * binary_tensor return output
def recurrent_func(theta: NestedMap, states_0: NestedMap, inputs: NestedMap, cell_fn: Callable[[NestedMap, NestedMap, NestedMap], NestedMap]): """Computes a recurrent neural net. Args: theta: weights. A `.NestedMap`. states_0: initial state. A `.NestedMap`. inputs: inputs. A `.NestedMap`. cell_fn: A python function which computes:: states_1 = cell_fn(theta, states_0, inputs[t, :]) Returns: `accumulate_state` and the final state. """ input_seq_len = inputs.Flatten()[0].shape[0] def assert_not_none(x): assert x is not None tf.nest.map_structure(assert_not_none, states_0) tf.nest.map_structure(assert_not_none, inputs) tf.nest.map_structure(assert_not_none, theta) def new_cum_state(x): x1 = jnp.expand_dims(x, 0) # +1 so that we can store initial_states at position 0. return jnp.tile(x1, [input_seq_len + 1] + [1] * x.ndim) cumulative_states = states_0.Transform(new_cum_state) prng_key = base_layer.next_prng_key() global_step = base_layer.cur_global_step() start_time = jnp.array(0, dtype=jnp.uint32) fwd_initial_loop_vars = NestedMap(cur_time=start_time, theta=theta, states_0=states_0, cumulative_states=cumulative_states, inputs=inputs) def same_type_shape(x, y): assert x.dtype == y.dtype, (x.dtype, y.dtype) assert x.shape == y.shape, (x.shape, y.shape) def wrapped_cell_fn(fn_in): # fn_in is NestedMap containing the following elements: # - t # - theta # - states_0 # - inputs_t # Start a chain of prng key that also takes into account of time steps. t = fn_in.t theta = fn_in.theta states_0 = fn_in.states_0 inputs_t = fn_in.inputs_t with base_layer.JaxContext.new_context(prng_key=jax.random.fold_in( prng_key, t), global_step=global_step): # NO side-effect ops are allowed as the enclosing JaxContext is not bound # to any layer. states_1 = cell_fn(theta, states_0, inputs_t) tf.nest.assert_same_structure(states_0, states_1) tf.nest.map_structure(same_type_shape, states_0, states_1) return states_1 def wrapped_cell_fn_grad(fn_in, d_fn_out): # This is roughly the following: # # fn_out = wrapped_cell_fn(fn_in) # d_fn_in = tf.gradient(fn_out, fn_in, d_fn_out) # return d_fn_in # assert isinstance(fn_in, NestedMap) fn_out, vjp_fn = jax.vjp(wrapped_cell_fn, fn_in) del fn_out d_fn_in = vjp_fn(d_fn_out) assert isinstance(d_fn_in, tuple) assert len(d_fn_in) == 1 d_fn_in_0 = d_fn_in[0] # Over-write gradient for t, the time step. d_fn_in_0.t = jnp.zeros_like(fn_in.t) tf.nest.assert_same_structure(fn_in, d_fn_in_0) tf.nest.map_structure(same_type_shape, fn_in, d_fn_in_0) return d_fn_in_0 def fwd_comp_fn(loop_vars): # loop_vars is a NestedMap containing the following elements: # - cur_time # - theta # - inputs # - cumulative_states # - states_0 t = loop_vars.cur_time theta = loop_vars.theta inputs = loop_vars.inputs cumulative_states = loop_vars.cumulative_states states_0 = loop_vars.states_0 inputs_t = inputs.Transform(lambda x: x[t]) states_1 = wrapped_cell_fn( NestedMap(t=t, theta=theta, states_0=states_0, inputs_t=inputs_t)) def set_t(x, x_t): return x.at[t + 1].set(x_t) cumulative_states = tf.nest.map_structure(set_t, cumulative_states, states_1) loop_out = NestedMap(cur_time=t + 1, theta=theta, inputs=inputs, states_0=states_1, cumulative_states=cumulative_states) return loop_out def fwd_continue_fn(loop_vars): return loop_vars.cur_time < input_seq_len # This custom_vjp implementation follows examples here: # https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html @jax.custom_vjp def fwd_loop(loop_vars): final_loop_vars = jax.lax.while_loop(fwd_continue_fn, fwd_comp_fn, loop_vars) return NestedMap(final_states=final_loop_vars.states_0, cumulative_states=final_loop_vars.cumulative_states) def loop_fn_vjp_fwd(loop_vars): loop_fn_out = fwd_loop(loop_vars) return loop_fn_out, (loop_vars, loop_fn_out.cumulative_states) def loop_fn_vjp_bwd(res, d_out): fwd_loop_vars, cumulative_states = res d_final_states = d_out.final_states d_cumulative_states = d_out.cumulative_states start_time = input_seq_len - 1 d_states_1 = tf.nest.map_structure(lambda x, y: x[start_time + 1] + y, d_cumulative_states, d_final_states) bwd_loop_vars = NestedMap( cur_time=start_time, theta=fwd_loop_vars.theta, inputs=fwd_loop_vars.inputs, cumulative_states=cumulative_states, d_cumulative_states=d_cumulative_states, d_theta=fwd_loop_vars.theta.Transform(jnp.zeros_like), d_inputs=fwd_loop_vars.inputs.Transform(jnp.zeros_like), d_states_1=d_states_1) def bwd_comp_fn(loop_vars): t = loop_vars.cur_time inputs = loop_vars.inputs inputs_t = inputs.Transform(lambda x: x[t]) states_0 = loop_vars.cumulative_states.Transform(lambda x: x[t]) d_cell_in = wrapped_cell_fn_grad( NestedMap(t=t, theta=loop_vars.theta, states_0=states_0, inputs_t=inputs_t), loop_vars.d_states_1) d_theta = tf.nest.map_structure(lambda x, y: x + y, loop_vars.d_theta, d_cell_in.theta) d_states_0 = tf.nest.map_structure(lambda x, y: x + y[t], d_cell_in.states_0, loop_vars.d_cumulative_states) def set_t(x, x_t): return x.at[t].set(x_t) d_inputs = tf.nest.map_structure(set_t, loop_vars.d_inputs, d_cell_in.inputs_t) loop_vars_out = loop_vars.Transform(lambda x: x) loop_vars_out.d_inputs = d_inputs loop_vars_out.d_states_1 = d_states_0 loop_vars_out.d_theta = d_theta loop_vars_out.cur_time = t - 1 return loop_vars_out def bwd_continue_fn(loop_vars): return loop_vars.cur_time >= 0 bwd_final_loop_vars = jax.lax.while_loop(bwd_continue_fn, bwd_comp_fn, bwd_loop_vars) d_out = fwd_loop_vars.Transform(jnp.zeros_like) tf.nest.map_structure(same_type_shape, d_out.states_0, bwd_final_loop_vars.d_states_1) tf.nest.map_structure(same_type_shape, d_out.theta, bwd_final_loop_vars.d_theta) tf.nest.map_structure(same_type_shape, d_out.inputs, bwd_final_loop_vars.d_inputs) d_out.states_0 = bwd_final_loop_vars.d_states_1 d_out.theta = bwd_final_loop_vars.d_theta d_out.inputs = bwd_final_loop_vars.d_inputs return (d_out, ) fwd_loop.defvjp(loop_fn_vjp_fwd, loop_fn_vjp_bwd) # Finally, let's simply run the forward loop fn. fwd_final_loop_vars = fwd_loop(fwd_initial_loop_vars) fwd_cumulative_states = fwd_final_loop_vars.cumulative_states.Transform( lambda x: x[1:]) return fwd_final_loop_vars.final_states, fwd_cumulative_states
def scan(carry_init: NestedMap, xs: NestedMap, fn: Callable[[NestedMap, NestedMap], Tuple[NestedMap, NestedMap]], root_layer: Optional[base_layer.BaseLayer] = None, checkpoint_policy: AutodiffCheckpointType = AutodiffCheckpointType. SAVE_NOTHING): """A simple wrap around jax.lax.scan. Back-prop is availale through auto-diff. Args: carry_init: initial state. A `.NestedMap`. xs: inputs. A `.NestedMap`. All inputs in time-major. fn: A python function which computes: carry, ys[t] = fn(carry, xs[t, :]) root_layer: The root layer within which this jax.lax.scan based while_loop is carried out. If root_layer is provided, some basic-effort check is performed to make sure fn is side-effect free. Otherwise, no such checks are performed. checkpoint_policy: A AutodiffCheckpointType. How to checkpoint for BProp: SAVE_NOTHING, SAVE_DOT_ONLY, SAVE_DOT_WITH_NO_BATCH_DIM. Returns: (final 'carry', 'ys', stacked summaries). """ del root_layer assert isinstance(carry_init, py_utils.NestedMap) assert isinstance(xs, py_utils.NestedMap) # Make a copy of carry_init structure. carry_init = tf.nest.map_structure(lambda x: x, carry_init) # "carry" will be augmented with the following three tensors, so make sure # they don't already exist in the NestedMap. assert 'time_step' not in carry_init assert 'prng_key' not in carry_init assert 'global_step' not in carry_init def custom_policy(checkpoint_policy: AutodiffCheckpointType): # TODO(zhangqiaorjc): Configure custom checkpoint policy in expt config # without introducing enum. if checkpoint_policy == AutodiffCheckpointType.SAVE_EVERYTHING: return jax.checkpoint_policies.everything_saveable if checkpoint_policy == AutodiffCheckpointType.SAVE_DOT_ONLY: return jax.checkpoint_policies.checkpoint_dots if checkpoint_policy == AutodiffCheckpointType.SAVE_DOT_WITH_NO_BATCH_DIM: return jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims if checkpoint_policy == AutodiffCheckpointType.SAVE_QKV_OUT_PROJ: return jax.checkpoint_policies.save_only_these_names( 'combined_qkv_proj', 'out_proj') if checkpoint_policy == AutodiffCheckpointType.SAVE_CONTEXT: return jax.checkpoint_policies.save_only_these_names('context') if checkpoint_policy == AutodiffCheckpointType.SAVE_OUT_PROJ: return jax.checkpoint_policies.save_only_these_names('out_proj') if checkpoint_policy == AutodiffCheckpointType.SAVE_CONTEXT_AND_OUT_PROJ: return jax.checkpoint_policies.save_only_these_names( 'context', 'out_proj') if checkpoint_policy == AutodiffCheckpointType.SAVE_DOT_FOR_MLPERF_200B: return jax.checkpoint_policies.save_only_these_names( 'combined_qkv_proj', 'query_proj', 'value_proj', 'key_proj', 'context', 'out_proj') assert checkpoint_policy == AutodiffCheckpointType.SAVE_NOTHING return jax.checkpoint_policies.nothing_saveable @functools.partial(ad_checkpoint.checkpoint, prevent_cse=False, policy=custom_policy(checkpoint_policy)) def fn_wrap(carry, xs_t): # carry is augmented with time_step, prng_key, global_step three additional # tensors to make fn_wrap fully functional. # Start a new prng_key branch that also depends on the time step. prng_key_t = jax.random.fold_in(carry.prng_key, carry.time_step) with base_layer.JaxContext.new_context(prng_key=prng_key_t, global_step=carry.global_step): carry_new, ys_t = fn(carry, xs_t) carry_new.time_step = carry.time_step + 1 # copy over prng_key and global_step carry_new.prng_key = carry.prng_key carry_new.global_step = carry.global_step tf.nest.assert_same_structure(carry_new, carry) summaries = base_layer.all_summaries() return carry_new, (ys_t, summaries) # The initial time step. time_step = jnp.array(0, dtype=jnp.uint32) prng_key = base_layer.next_prng_key() global_step = base_layer.cur_global_step() carry_init.time_step = time_step carry_init.prng_key = prng_key carry_init.global_step = global_step carry_final, (ys, summaries) = jax.lax.scan(fn_wrap, carry_init, xs) del carry_final.time_step del carry_final.global_step del carry_final.prng_key return carry_final, ys, summaries
def recurrent_static(theta: NestedMap, states_0: NestedMap, inputs: NestedMap, cell_fn: Callable[[NestedMap, NestedMap, NestedMap], NestedMap], root_layer: Optional[base_layer.BaseLayer] = None): """A simpler form of Recurrent where num of steps is known statically. Back-prop is availale through auto-diff. 'padding' in inputs is used to skip certain steps dynamically. If the 'padding' tensor exists, it is expected of a binary 0/1 tensor. Args: theta: weights. A `.NestedMap`. states_0: initial state. A `.NestedMap`. inputs: inputs. A `.NestedMap`. All inputs in time-major. cell_fn: A python function which computes:: states_1 = cell_fn(theta, states_0, inputs[t, :]) root_layer: The root layer within which this recurrent_static recurrent loop is carried out. Returns: `accumulate_state` and the final state. """ assert 'time_step' not in states_0 # The initial time step. time_step = jnp.array(0, dtype=jnp.uint32) # Make a copy of states_0 structure. states_0 = tf.nest.map_structure(lambda x: x, states_0) states_0.time_step = time_step prng_key = base_layer.next_prng_key() global_step = base_layer.cur_global_step() # TODO(zhangqiaorjc): Switch to ad_checkpoint.checkpoint after mattjj bug fix. @jax.checkpoint def comp_fn(states_0, inputs_t): # Start a new prng_key branch that also depends on the time step. if root_layer is not None: forward_updated_vars_before = tf.nest.map_structure( lambda x: x, root_layer.forward_updated_vars) prng_key_t = jax.random.fold_in(prng_key, states_0.time_step) with base_layer.JaxContext.new_context(prng_key=prng_key_t, global_step=global_step): # NO side-effect ops are allowed as the enclosing JaxContext is not bound # to any layer. # # Whether or not we should skip this time step. if 'padding' in inputs_t: # We skip if all are padded steps. skip = jnp.all(inputs_t.padding > 0.5) else: skip = jnp.array(False) def carry_over(args): states_0, inputs_t = args del inputs_t # We simply carry over the states for this time step. states_1 = tf.nest.map_structure(lambda x: x, states_0) states_1.time_step = states_0.time_step + 1 return states_1 def do_compute(args): states_0, inputs_t = args # Actually carry out the computation. states_1 = cell_fn(theta, states_0, inputs_t) states_1.time_step = states_0.time_step + 1 return states_1 if 'padding' in inputs_t: states_1 = jax.lax.cond(skip, carry_over, do_compute, (states_0, inputs_t)) else: states_1 = do_compute((states_0, inputs_t)) tf.nest.assert_same_structure(states_0, states_1) if root_layer is not None: forward_updated_vars_after = tf.nest.map_structure( lambda x: x, root_layer.forward_updated_vars) def assert_no_change(x, y): assert (x is None and y is None) or (x is not None and y is not None) tf.nest.map_structure(assert_no_change, forward_updated_vars_before, forward_updated_vars_after) return states_1, states_1 final_states, cumulative_states = jax.lax.scan(comp_fn, states_0, inputs) del final_states.time_step del cumulative_states.time_step return final_states, cumulative_states
def _UniformSample(sample_p: float) -> JTensor: prng_key = base_layer.next_prng_key() rnd_sample = jax.random.uniform(prng_key, inputs.shape) return (rnd_sample < sample_p).astype(fprop_dtype)
def fprop(self, inputs: JTensor, paddings: JTensor) -> Tuple[JTensor, JTensor]: """Applies data augmentation by randomly masking/replacing tokens in inputs. Args: inputs: An int32 tensor of shape [batch, length]. paddings: A 0/1 tensor of shape [batch, length]. Returns: A pair <new_inputs, mask>: new_inputs: An int32 tensor of shape [batch, length]. The new token ids after data augmentation. mask: A 0/1 tensor. A "1" indicates the corresponding token at that position had undergone the data augmentation process. """ p = self.params assert p.vocab_size > 0 assert p.mask_token_id >= 0 assert p.mask_prob + p.random_prob + p.same_prob < 1.0 assert p.mask_prob + p.random_prob + p.same_prob > 0.0 fprop_dtype = self.fprop_dtype def _UniformSample(sample_p: float) -> JTensor: prng_key = base_layer.next_prng_key() rnd_sample = jax.random.uniform(prng_key, inputs.shape) return (rnd_sample < sample_p).astype(fprop_dtype) total_replacement_prob = p.mask_prob + p.random_prob + p.same_prob # valid_tokens == 1.0 if the corresponding position is a valid token. valid_tokens = 1.0 - paddings.astype(fprop_dtype) # replacement == 1.0 if the corresponding token is to be replaced by # something else (mask, random, self). replacement_pos = valid_tokens * _UniformSample(total_replacement_prob) no_replacement = 1.0 - replacement_pos # First sample the token positions to be masked out. remaining_prob = total_replacement_prob remaining_pos = replacement_pos mask_prob = p.mask_prob / remaining_prob # mask_pos == 1.0 if the corresponding token should be masked. mask_pos = remaining_pos * _UniformSample(mask_prob) # Next sample the token positions to be replaced by random tokens. remaining_prob -= p.mask_prob remaining_pos -= mask_pos assert remaining_prob > 0.0 random_prob = p.random_prob / remaining_prob random_pos = remaining_pos * _UniformSample(random_prob) # Lastly, token positions to be replaced by self. self_pos = remaining_pos - random_pos random_tokens = jax.random.randint(base_layer.next_prng_key(), inputs.shape, 0, p.vocab_size, inputs.dtype) mask_tokens = jnp.zeros_like(inputs) + p.mask_token_id input_dtype = inputs.dtype augmented = (inputs * no_replacement.astype(input_dtype) + mask_tokens * mask_pos.astype(input_dtype) + random_tokens * random_pos.astype(input_dtype) + inputs * self_pos.astype(input_dtype)) return augmented, replacement_pos