Skip to content

Commit bec207a

Browse files
dxqbdxqb
committed
Merge PR Nerogar#1476 (Make gradient checkpointing and offloading per-component) into preview
2 parents 97a2fc2 + 5a41835 commit bec207a

48 files changed

Lines changed: 452 additions & 321 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎modules/model/ChromaModel.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -113,8 +113,7 @@ def vae_to(self, device: torch.device):
113113

114114
def text_encoder_to(self, device: torch.device):
115115
if self.text_encoder is not None:
116-
if self.text_encoder_offload_conductor is not None and \
117-
self.text_encoder_offload_conductor.layer_offload_activated():
116+
if self.text_encoder_offload_conductor is not None:
118117
self.text_encoder_offload_conductor.to(device)
119118
else:
120119
self.text_encoder.to(device=device)
@@ -123,8 +122,7 @@ def text_encoder_to(self, device: torch.device):
123122
self.text_encoder_lora.to(device)
124123

125124
def transformer_to(self, device: torch.device):
126-
if self.transformer_offload_conductor is not None and \
127-
self.transformer_offload_conductor.layer_offload_activated():
125+
if self.transformer_offload_conductor is not None:
128126
self.transformer_offload_conductor.to(device)
129127
else:
130128
self.transformer.to(device=device)

‎modules/model/ErnieModel.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,15 +72,13 @@ def vae_to(self, device: torch.device):
7272

7373
def text_encoder_to(self, device: torch.device):
7474
if self.text_encoder is not None:
75-
if self.text_encoder_offload_conductor is not None and \
76-
self.text_encoder_offload_conductor.layer_offload_activated():
75+
if self.text_encoder_offload_conductor is not None:
7776
self.text_encoder_offload_conductor.to(device)
7877
else:
7978
self.text_encoder.to(device=device)
8079

8180
def transformer_to(self, device: torch.device):
82-
if self.transformer_offload_conductor is not None and \
83-
self.transformer_offload_conductor.layer_offload_activated():
81+
if self.transformer_offload_conductor is not None:
8482
self.transformer_offload_conductor.to(device)
8583
else:
8684
self.transformer.to(device=device)

‎modules/model/Flux2Model.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -122,15 +122,13 @@ def vae_to(self, device: torch.device):
122122

123123
def text_encoder_to(self, device: torch.device):
124124
if self.text_encoder is not None:
125-
if self.text_encoder_offload_conductor is not None and \
126-
self.text_encoder_offload_conductor.layer_offload_activated():
125+
if self.text_encoder_offload_conductor is not None:
127126
self.text_encoder_offload_conductor.to(device)
128127
else:
129128
self.text_encoder.to(device=device)
130129

131130
def transformer_to(self, device: torch.device):
132-
if self.transformer_offload_conductor is not None and \
133-
self.transformer_offload_conductor.layer_offload_activated():
131+
if self.transformer_offload_conductor is not None:
134132
self.transformer_offload_conductor.to(device)
135133
else:
136134
self.transformer.to(device=device)

‎modules/model/FluxModel.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -149,8 +149,7 @@ def text_encoder_1_to(self, device: torch.device):
149149

150150
def text_encoder_2_to(self, device: torch.device):
151151
if self.text_encoder_2 is not None:
152-
if self.text_encoder_2_offload_conductor is not None and \
153-
self.text_encoder_2_offload_conductor.layer_offload_activated():
152+
if self.text_encoder_2_offload_conductor is not None:
154153
self.text_encoder_2_offload_conductor.to(device)
155154
else:
156155
self.text_encoder_2.to(device=device)
@@ -159,8 +158,7 @@ def text_encoder_2_to(self, device: torch.device):
159158
self.text_encoder_2_lora.to(device)
160159

161160
def transformer_to(self, device: torch.device):
162-
if self.transformer_offload_conductor is not None and \
163-
self.transformer_offload_conductor.layer_offload_activated():
161+
if self.transformer_offload_conductor is not None:
164162
self.transformer_offload_conductor.to(device)
165163
else:
166164
self.transformer.to(device=device)

‎modules/model/HiDreamModel.py‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -220,8 +220,7 @@ def text_encoder_2_to(self, device: torch.device):
220220

221221
def text_encoder_3_to(self, device: torch.device):
222222
if self.text_encoder_3 is not None:
223-
if self.text_encoder_3_offload_conductor is not None and \
224-
self.text_encoder_3_offload_conductor.layer_offload_activated():
223+
if self.text_encoder_3_offload_conductor is not None:
225224
self.text_encoder_3_offload_conductor.to(device)
226225
else:
227226
self.text_encoder_3.to(device=device)
@@ -231,8 +230,7 @@ def text_encoder_3_to(self, device: torch.device):
231230

232231
def text_encoder_4_to(self, device: torch.device):
233232
if self.text_encoder_4 is not None:
234-
if self.text_encoder_4_offload_conductor is not None and \
235-
self.text_encoder_4_offload_conductor.layer_offload_activated():
233+
if self.text_encoder_4_offload_conductor is not None:
236234
self.text_encoder_4_offload_conductor.to(device)
237235
else:
238236
self.text_encoder_4.to(device=device)
@@ -241,8 +239,7 @@ def text_encoder_4_to(self, device: torch.device):
241239
self.text_encoder_4_lora.to(device)
242240

243241
def transformer_to(self, device: torch.device):
244-
if self.transformer_offload_conductor is not None and \
245-
self.transformer_offload_conductor.layer_offload_activated():
242+
if self.transformer_offload_conductor is not None:
246243
self.transformer_offload_conductor.to(device)
247244
else:
248245
self.transformer.to(device=device)

‎modules/model/HunyuanVideoModel.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -157,8 +157,7 @@ def text_encoder_to(self, device: torch.device):
157157

158158
def text_encoder_1_to(self, device: torch.device):
159159
if self.text_encoder_1 is not None:
160-
if self.text_encoder_1_offload_conductor is not None and \
161-
self.text_encoder_1_offload_conductor.layer_offload_activated():
160+
if self.text_encoder_1_offload_conductor is not None:
162161
self.text_encoder_1_offload_conductor.to(device)
163162
else:
164163
self.text_encoder_1.to(device=device)
@@ -174,8 +173,7 @@ def text_encoder_2_to(self, device: torch.device):
174173
self.text_encoder_2_lora.to(device)
175174

176175
def transformer_to(self, device: torch.device):
177-
if self.transformer_offload_conductor is not None and \
178-
self.transformer_offload_conductor.layer_offload_activated():
176+
if self.transformer_offload_conductor is not None:
179177
self.transformer_offload_conductor.to(device)
180178
else:
181179
self.transformer.to(device=device)

‎modules/model/PixArtAlphaModel.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,7 @@ def vae_to(self, device: torch.device):
114114
self.vae.to(device=device)
115115

116116
def text_encoder_to(self, device: torch.device):
117-
if self.text_encoder_offload_conductor is not None and \
118-
self.text_encoder_offload_conductor.layer_offload_activated():
117+
if self.text_encoder_offload_conductor is not None:
119118
self.text_encoder_offload_conductor.to(device)
120119
else:
121120
self.text_encoder.to(device=device)
@@ -124,8 +123,7 @@ def text_encoder_to(self, device: torch.device):
124123
self.text_encoder_lora.to(device)
125124

126125
def transformer_to(self, device: torch.device):
127-
if self.transformer_offload_conductor is not None and \
128-
self.transformer_offload_conductor.layer_offload_activated():
126+
if self.transformer_offload_conductor is not None:
129127
self.transformer_offload_conductor.to(device)
130128
else:
131129
self.transformer.to(device=device)

‎modules/model/QwenModel.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,7 @@ def vae_to(self, device: torch.device):
8181

8282
def text_encoder_to(self, device: torch.device): #TODO share more code between models
8383
if self.text_encoder is not None:
84-
if self.text_encoder_offload_conductor is not None and \
85-
self.text_encoder_offload_conductor.layer_offload_activated():
84+
if self.text_encoder_offload_conductor is not None:
8685
self.text_encoder_offload_conductor.to(device)
8786
else:
8887
self.text_encoder.to(device=device)
@@ -91,8 +90,7 @@ def text_encoder_to(self, device: torch.device): #TODO share more code between m
9190
self.text_encoder_lora.to(device)
9291

9392
def transformer_to(self, device: torch.device):
94-
if self.transformer_offload_conductor is not None and \
95-
self.transformer_offload_conductor.layer_offload_activated():
93+
if self.transformer_offload_conductor is not None:
9694
self.transformer_offload_conductor.to(device)
9795
else:
9896
self.transformer.to(device=device)

‎modules/model/SanaModel.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -116,8 +116,7 @@ def vae_to(self, device: torch.device):
116116
self.vae.to(device=device)
117117

118118
def text_encoder_to(self, device: torch.device):
119-
if self.text_encoder_offload_conductor is not None and \
120-
self.text_encoder_offload_conductor.layer_offload_activated():
119+
if self.text_encoder_offload_conductor is not None:
121120
self.text_encoder_offload_conductor.to(device)
122121
else:
123122
self.text_encoder.to(device=device)
@@ -126,8 +125,7 @@ def text_encoder_to(self, device: torch.device):
126125
self.text_encoder_lora.to(device)
127126

128127
def transformer_to(self, device: torch.device):
129-
if self.transformer_offload_conductor is not None and \
130-
self.transformer_offload_conductor.layer_offload_activated():
128+
if self.transformer_offload_conductor is not None:
131129
self.transformer_offload_conductor.to(device)
132130
else:
133131
self.transformer.to(device=device)

‎modules/model/StableDiffusion3Model.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,7 @@ def text_encoder_2_to(self, device: torch.device):
180180

181181
def text_encoder_3_to(self, device: torch.device):
182182
if self.text_encoder_3 is not None:
183-
if self.text_encoder_3_offload_conductor is not None and \
184-
self.text_encoder_3_offload_conductor.layer_offload_activated():
183+
if self.text_encoder_3_offload_conductor is not None:
185184
self.text_encoder_3_offload_conductor.to(device)
186185
else:
187186
self.text_encoder_3.to(device=device)
@@ -190,8 +189,7 @@ def text_encoder_3_to(self, device: torch.device):
190189
self.text_encoder_3_lora.to(device)
191190

192191
def transformer_to(self, device: torch.device):
193-
if self.transformer_offload_conductor is not None and \
194-
self.transformer_offload_conductor.layer_offload_activated():
192+
if self.transformer_offload_conductor is not None:
195193
self.transformer_offload_conductor.to(device)
196194
else:
197195
self.transformer.to(device=device)

0 commit comments

Comments
 (0)