rxnn 0.1.62__py3-none-any.whl → 0.1.64__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- rxnn/rxt/models.py +17 -13
- rxnn/training/base.py +6 -2
- {rxnn-0.1.62.dist-info → rxnn-0.1.64.dist-info}/METADATA +1 -1
- {rxnn-0.1.62.dist-info → rxnn-0.1.64.dist-info}/RECORD +6 -6
- {rxnn-0.1.62.dist-info → rxnn-0.1.64.dist-info}/LICENSE +0 -0
- {rxnn-0.1.62.dist-info → rxnn-0.1.64.dist-info}/WHEEL +0 -0
rxnn/rxt/models.py
CHANGED
|
@@ -31,9 +31,11 @@ class RxTAlphaComponentConfig(TypedDict):
|
|
|
31
31
|
moe_top_k: int
|
|
32
32
|
self_att_type: str
|
|
33
33
|
cross_att_type: str
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
34
|
+
att_experts: int
|
|
35
|
+
att_query_experts: int
|
|
36
|
+
att_query_groups: int
|
|
37
|
+
cross_att_groups: int
|
|
38
|
+
cross_att_query_groups: int
|
|
37
39
|
|
|
38
40
|
|
|
39
41
|
class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
|
|
@@ -61,9 +63,11 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
|
|
|
61
63
|
moe_top_k: int = 1,
|
|
62
64
|
self_att_type: str = 'gqa',
|
|
63
65
|
cross_att_type: str = 'mqa',
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
66
|
+
att_experts: int = None,
|
|
67
|
+
att_query_experts: int = None,
|
|
68
|
+
att_query_groups: int = None,
|
|
69
|
+
cross_att_groups: int = None,
|
|
70
|
+
cross_att_query_groups: int = None,
|
|
67
71
|
**kwargs
|
|
68
72
|
):
|
|
69
73
|
super(RxTAlphaComponentBase, self).__init__(**kwargs)
|
|
@@ -86,20 +90,20 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
|
|
|
86
90
|
else:
|
|
87
91
|
att_init = lambda: init_experimental_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
|
|
88
92
|
use_flash_attention=use_flash_attention, dropout=att_dropout,
|
|
89
|
-
max_seq_len=seq_len, is_causal=is_causal, num_experts=
|
|
90
|
-
num_query_experts=
|
|
91
|
-
num_query_groups=
|
|
93
|
+
max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
|
|
94
|
+
num_query_experts=att_query_experts,
|
|
95
|
+
num_query_groups=att_query_groups)
|
|
92
96
|
|
|
93
97
|
if cross_att_type in ['mha', 'gqa', 'mqa']:
|
|
94
98
|
cross_att_init = lambda: init_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
|
|
95
99
|
use_flash_attention=use_flash_attention, dropout=att_dropout,
|
|
96
100
|
max_seq_len=seq_len, is_causal=is_causal)
|
|
97
101
|
else:
|
|
98
|
-
cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
|
|
102
|
+
cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type, cross_att_groups or att_groups, rope=rope,
|
|
99
103
|
use_flash_attention=use_flash_attention, dropout=att_dropout,
|
|
100
|
-
max_seq_len=seq_len, is_causal=is_causal, num_experts=
|
|
101
|
-
num_query_experts=
|
|
102
|
-
num_query_groups=
|
|
104
|
+
max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
|
|
105
|
+
num_query_experts=att_query_experts,
|
|
106
|
+
num_query_groups=cross_att_query_groups or att_query_groups)
|
|
103
107
|
|
|
104
108
|
layers = nn.ModuleList([
|
|
105
109
|
ReactiveTransformerLayer(
|
rxnn/training/base.py
CHANGED
|
@@ -113,7 +113,9 @@ class BaseTrainer(ABC):
|
|
|
113
113
|
|
|
114
114
|
self.model.train()
|
|
115
115
|
for epoch in range(self.current_epoch, self.current_epoch + epochs):
|
|
116
|
-
if self.is_running:
|
|
116
|
+
if not self.is_running:
|
|
117
|
+
break
|
|
118
|
+
else:
|
|
117
119
|
self.current_epoch = epoch
|
|
118
120
|
self.epoch_steps = 0
|
|
119
121
|
if train_sampler is not None:
|
|
@@ -143,7 +145,9 @@ class BaseTrainer(ABC):
|
|
|
143
145
|
self.optimizer_step_count = 0
|
|
144
146
|
|
|
145
147
|
for batch_idx, batch in enumerate(dataloader):
|
|
146
|
-
if self.is_running:
|
|
148
|
+
if not self.is_running:
|
|
149
|
+
break
|
|
150
|
+
else:
|
|
147
151
|
for callback in self.callbacks:
|
|
148
152
|
callback.on_batch_start(self.model, batch_idx, batch)
|
|
149
153
|
if self.get_batch_size(batch) == batch_size:
|
|
@@ -7,9 +7,9 @@ rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
|
7
7
|
rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
|
|
8
8
|
rxnn/memory/stm.py,sha256=EsD8slSP4_9dLuq6aFPDmuFe8PWilxh90so5Z3nm-ig,2057
|
|
9
9
|
rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
10
|
-
rxnn/rxt/models.py,sha256=
|
|
10
|
+
rxnn/rxt/models.py,sha256=9xJfb1rH7-QVO6PRsvUcbhskb1K7JTcE2ChwR4qT4EY,8711
|
|
11
11
|
rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
12
|
-
rxnn/training/base.py,sha256=
|
|
12
|
+
rxnn/training/base.py,sha256=xPMA2Bg9-oUZvSZg67ls2p7Gk9pZ9IHUiIJwUzSe2K8,11766
|
|
13
13
|
rxnn/training/bml.py,sha256=S1ZaXTybzeJH7uVFamCr4TPl2bLyZ5xmn_lSsjThTiM,19162
|
|
14
14
|
rxnn/training/callbacks.py,sha256=_YfMKY_eFdc-tubhO9nYH2PXDZDQwlSI74FVOoCXpQg,22108
|
|
15
15
|
rxnn/training/dataset.py,sha256=JQuWSUdT5AnsrG6M_EsewoU6uroVHhg4K715nbtDx8A,9643
|
|
@@ -25,7 +25,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
|
|
|
25
25
|
rxnn/transformers/positional.py,sha256=2l38RS0Dini3f6Z3LUHr3XwWzg1UK7fO2C6wazWDAYU,4292
|
|
26
26
|
rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
|
|
27
27
|
rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
|
|
28
|
-
rxnn-0.1.
|
|
29
|
-
rxnn-0.1.
|
|
30
|
-
rxnn-0.1.
|
|
31
|
-
rxnn-0.1.
|
|
28
|
+
rxnn-0.1.64.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
|
|
29
|
+
rxnn-0.1.64.dist-info/METADATA,sha256=Yyt5z1luYJbSbj41Bi1tZOkmRvcrF8p4rNUzO_alNWM,16579
|
|
30
|
+
rxnn-0.1.64.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
|
|
31
|
+
rxnn-0.1.64.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|