easydel.modules.rwkv.__init__

easydel.modules.rwkv.__init__#

class easydel.modules.rwkv.__init__.FlaxRwkvForCausalLM(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)[source]#

Bases: Module

config: RwkvConfig#
dtype#

alias of float32

name: Optional[str] = None#
param_dtype#

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None#
precision: Optional[Union[str, Precision]] = None#
scope: Scope | None = None#
setup()[source]#

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

  1. Immediately when invoking apply(), init() or init_and_output().

  2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

    >>> class MyModule(nn.Module):
...   def setup(self):
...     submodule = nn.Conv(...)

...     # Accessing `submodule` attributes does not yet work here.

...     # The following line invokes `self.__setattr__`, which gives
...     # `submodule` the name "conv1".
...     self.conv1 = submodule

...     # Accessing `submodule` attributes or methods is now safe and
...     # either causes setup() to be called once.

  3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class easydel.modules.rwkv.__init__.FlaxRwkvModel(config: easydel.modules.rwkv.rwkv_configuration.RwkvConfig, dtype: numpy.dtype = <class 'jax.numpy.float32'>, param_dtype: numpy.dtype = <class 'jax.numpy.float32'>, precision: Union[str, jax._src.lax.lax.Precision, NoneType] = None, parent: Union[flax.linen.module.Module, flax.core.scope.Scope, flax.linen.module._Sentinel, NoneType] = <flax.linen.module._Sentinel object at 0x7ec913f46950>, name: Optional[str] = None)[source]#

Bases: Module

config: RwkvConfig#
dtype#

alias of float32

name: Optional[str] = None#
param_dtype#

alias of float32

parent: Optional[Union[Module, Scope, _Sentinel]] = None#
precision: Optional[Union[str, Precision]] = None#
scope: Scope | None = None#
setup()[source]#

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

  1. Immediately when invoking apply(), init() or init_and_output().

  2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

    >>> class MyModule(nn.Module):
...   def setup(self):
...     submodule = nn.Conv(...)

...     # Accessing `submodule` attributes does not yet work here.

...     # The following line invokes `self.__setattr__`, which gives
...     # `submodule` the name "conv1".
...     self.conv1 = submodule

...     # Accessing `submodule` attributes or methods is now safe and
...     # either causes setup() to be called once.

  3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class easydel.modules.rwkv.__init__.RwkvConfig(vocab_size=50277, context_length=1024, hidden_size=4096, num_hidden_layers=32, attention_hidden_size=None, intermediate_size=None, layer_norm_epsilon=1e-05, bos_token_id=0, eos_token_id=0, rescale_every=6, tie_word_embeddings=False, use_cache=True, bits: Optional[int] = None, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, **kwargs)[source]#

Bases: EasyDeLBaseConfig

Configuration objects inherit from [EasyDeLBaseConfig] and can be used to control the model outputs. Read the documentation from [EasyDeLBaseConfig] for more information.

Parameters

  • vocab_size (int, optional, defaults to 50277) – Vocabulary size of the RWKV model. Defines the number of different tokens that can be represented by the inputs_ids passed when calling [~easydel.modules.RwkvModel].

  • context_length (int, optional, defaults to 1024) – The maximum sequence length that this model might ever be used with.

  • hidden_size (int, optional, defaults to 4096) – Dimensionality of the encoder layers and the pooler layer.

  • num_hidden_layers (int, optional, defaults to 32) – Number of hidden layers in the Transformer encoder.

  • attention_hidden_size (int, optional) – Dimensionality of the query/key/value of the MultiHead Attention layer of the RWKV* model. If None, it is set to hidden_size.

  • intermediate_size (int, optional) – Dimensionality of the “intermediate” (often named feed-forward) layer in the Transformer encoder. If None, it is set to 4 * hidden_size.

  • layer_norm_epsilon (float, optional, defaults to 1e-5) – The epsilon used by the layer normalization layers.

  • rescale_every (int, optional, defaults to 6) – Interval of layers at which to rescale the attention scores.

  • use_cache (bool, optional, defaults to True) – Whether or not the model should return the last key/values attentions (not used by all models).

  • bos_token_id (int, optional, defaults to 0) – The id for the beginning of stream token.

  • eos_token_id (int, optional, defaults to 0) – The id for the end of stream token.

  • tie_word_embeddings (bool, optional, defaults to False) – Whether to tie the weights of the input embeddings and the output embeddings.

  • bits (int, optional) – The number of bits to quantize the model to. If None, the model is not quantized.

  • gradient_checkpointing (str, optional, defaults to “nothing_saveable”) – What to save during gradient checkpointing. Choose one of “nothing_saveable”, “first_half_saveable”, “full_saveable”.

add_jax_args(bits: Optional[int] = None, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, **kwargs)[source]#
attribute_map: Dict[str, str] = {'max_position_embeddings': 'context_length'}#
get_partition_rules(fully_sharded_data_parallel: bool = True)[source]#

Get the partition rules for the model.

Parameters

fully_sharded_data_parallel (bool, optional, defaults to True) – Whether to use fully sharded data parallelism.

Returns

The partition rules.

Return type

tp.Tuple[tp.Tuple[str, PartitionSpec]]

model_type: str = 'rwkv'#
Contents