Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
ColossalAI
Commits
d4fb7bfd
Unverified
Commit
d4fb7bfd
authored
Jun 15, 2023
by
digger yu
Committed by
GitHub
Jun 15, 2023
Browse files
fix typo applications/Chat/coati/ (#3947)
parent
e8ad3c88
Changes
3
Hide whitespace changes
Inline
Side-by-side
Showing
3 changed files
with
12 additions
and
12 deletions
+12
-12
applications/Chat/coati/ray/experience_maker_holder.py
applications/Chat/coati/ray/experience_maker_holder.py
+4
-4
applications/Chat/coati/ray/lora_constructor.py
applications/Chat/coati/ray/lora_constructor.py
+6
-6
applications/Chat/coati/trainer/strategies/colossalai.py
applications/Chat/coati/trainer/strategies/colossalai.py
+2
-2
No files found.
applications/Chat/coati/ray/experience_maker_holder.py
View file @
d4fb7bfd
...
...
@@ -205,15 +205,15 @@ class ExperienceMakerHolder:
self
.
experience_maker
.
actor
.
model
.
load_state_dict
(
new_actor_state_dict
,
strict
=
False
)
else
:
new_actor_state_dict
=
state_dict_to
(
new_actor_state_dict
,
device
=
torch
.
cuda
.
current_device
())
state_dict_increas
a
e
=
self
.
actor_lora_constructor
.
reconstruct_increase
(
new_actor_state_dict
,
new_actor_lora_config_dict
)
self
.
actor_lora_constructor
.
load_state_dict_increase
(
self
.
experience_maker
.
actor
.
model
,
state_dict_increas
a
e
)
state_dict_increase
=
self
.
actor_lora_constructor
.
reconstruct_increase
(
new_actor_state_dict
,
new_actor_lora_config_dict
)
self
.
actor_lora_constructor
.
load_state_dict_increase
(
self
.
experience_maker
.
actor
.
model
,
state_dict_increase
)
if
new_critic_state_dict
is
not
None
:
if
not
self
.
_update_lora_weights
or
fully_update
:
self
.
experience_maker
.
critic
.
load_state_dict
(
new_critic_state_dict
,
strict
=
False
)
else
:
new_critic_state_dict
=
state_dict_to
(
new_critic_state_dict
,
device
=
torch
.
cuda
.
current_device
())
state_dict_increas
a
e
=
self
.
critic_lora_constructor
.
reconstruct_increase
(
new_critic_state_dict
,
new_critic_lora_config_dict
)
self
.
critic_lora_constructor
.
load_state_dict_increase
(
self
.
experience_maker
.
critic
,
state_dict_increas
a
e
)
state_dict_increase
=
self
.
critic_lora_constructor
.
reconstruct_increase
(
new_critic_state_dict
,
new_critic_lora_config_dict
)
self
.
critic_lora_constructor
.
load_state_dict_increase
(
self
.
experience_maker
.
critic
,
state_dict_increase
)
# the lock must be released after both actor and critic being updated
if
chunk_end
:
...
...
applications/Chat/coati/ray/lora_constructor.py
View file @
d4fb7bfd
...
...
@@ -19,7 +19,7 @@ class LoRAConfig:
class
LoRAConstructor
:
'''
Tools for reconstructing a model from a remote LoRA model.
(Transfering only LoRA data costs much less!)
(Transfer
r
ing only LoRA data costs much less!)
Usage:
Step 1 (Sender):
filter_state_dict_lora()
...
...
@@ -52,7 +52,7 @@ class LoRAConstructor:
if
lora_config_dict
is
not
None
:
self
.
register_lora_config
(
lora_config_dict
)
state_dict_increas
a
e
=
OrderedDict
()
state_dict_increase
=
OrderedDict
()
config_iter
=
iter
(
self
.
lora_config_dict
.
items
())
lora_A
,
lora_B
,
layer_prefix
=
None
,
None
,
None
for
k
,
v
in
state_dict_lora
.
items
():
...
...
@@ -65,11 +65,11 @@ class LoRAConstructor:
assert
layer_prefix_2
==
layer_prefix
,
"unmatched (state_dict, config_dict) pair"
lora_B
=
v
weight_data_increase
=
self
.
_compute
(
lora_A
,
lora_B
,
config
)
state_dict_increas
a
e
[
layer_prefix
+
'.weight'
]
=
weight_data_increase
state_dict_increase
[
layer_prefix
+
'.weight'
]
=
weight_data_increase
lora_A
,
lora_B
,
layer_prefix
=
None
,
None
,
None
else
:
raise
ValueError
(
'unexpected key'
)
return
state_dict_increas
a
e
return
state_dict_increase
def
_compute
(
self
,
lora_A
,
lora_B
,
config
=
LoRAConfig
()):
def
T
(
w
):
...
...
@@ -80,12 +80,12 @@ class LoRAConstructor:
return
weight_data_increase
return
0
def
load_state_dict_increase
(
self
,
model
:
nn
.
Module
,
state_dict_increas
a
e
:
Dict
[
str
,
Any
]):
def
load_state_dict_increase
(
self
,
model
:
nn
.
Module
,
state_dict_increase
:
Dict
[
str
,
Any
]):
'''
The final reconstruction step
'''
# naive approach
model
.
load_state_dict
({
k
:
v
+
model
.
state_dict
()[
k
]
for
k
,
v
in
state_dict_increas
a
e
.
items
()},
strict
=
False
)
model
.
load_state_dict
({
k
:
v
+
model
.
state_dict
()[
k
]
for
k
,
v
in
state_dict_increase
.
items
()},
strict
=
False
)
@
staticmethod
def
filter_state_dict_lora
(
state_dict
:
Dict
[
str
,
Any
],
keep_non_lora
=
False
):
...
...
applications/Chat/coati/trainer/strategies/colossalai.py
View file @
d4fb7bfd
...
...
@@ -29,7 +29,7 @@ class ColossalAIStrategy(DDPStrategy):
precision(str): The precision to use. Choose in ('fp32', 'fp16'). Stage 3 only supports fp16.
seed(int): The seed for the random number generator.
shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3.
This is not compati
v
le with `from_pretrained()`. We temporarily disable this and will support it in the future.
This is not compati
b
le with `from_pretrained()`. We temporarily disable this and will support it in the future.
placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda')
If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU,
If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest.
...
...
@@ -39,7 +39,7 @@ class ColossalAIStrategy(DDPStrategy):
hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3.
min_chunk_size_mb(float): The minimum chunk size in MB. Only for ZeRO-3.
gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3.
reduce_bu
g
ket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2.
reduce_bu
c
ket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2.
overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2.
initial_scale(float): The initial scale for the optimizer.
growth_factor(float): The growth factor for the optimizer.
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment