@@ -103,8 +103,8 @@ def __init__(
103
103
tasks : str ,
104
104
pipeline_parameters : PipelineParameters ,
105
105
evaluation_tracker : EvaluationTracker ,
106
- model = None ,
107
106
model_config = None ,
107
+ model = None ,
108
108
):
109
109
if not (model or model_config ):
110
110
raise ValueError ("Must provide either a model or model config when creating a pipeline." )
@@ -116,10 +116,9 @@ def __init__(
116
116
"WARNING: --max_samples WAS SET. THESE NUMBERS ARE ONLY PARTIAL AND SHOULD NOT BE USED FOR COMPARISON UNLESS YOU KNOW WHAT YOU ARE DOING."
117
117
)
118
118
119
- self .accelerator , self .parallel_context = self ._init_parallelism_manager ()
120
-
121
- self .evaluation_tracker = evaluation_tracker
122
119
self .model_config = model_config
120
+ self .evaluation_tracker = evaluation_tracker
121
+ self .accelerator , self .parallel_context = self ._init_parallelism_manager ()
123
122
self .model = self ._init_model (model_config , model )
124
123
125
124
self .evaluation_tracker .general_config_logger .log_model_info (self .model .model_info )
@@ -141,9 +140,9 @@ def _init_parallelism_manager(self):
141
140
raise ValueError ("You are trying to launch a nanotron model, but nanotron is not installed" )
142
141
dist .initialize_torch_distributed ()
143
142
parallel_context = ParallelContext (
144
- tensor_parallel_size = self .model_config .parallelism .tp ,
145
- pipeline_parallel_size = self .model_config .parallelism .pp ,
146
- data_parallel_size = self .model_config .parallelism .dp ,
143
+ tensor_parallel_size = self .model_config .lighteval . parallelism .tp ,
144
+ pipeline_parallel_size = self .model_config .lighteval . parallelism .pp ,
145
+ data_parallel_size = self .model_config .lighteval . parallelism .dp ,
147
146
)
148
147
test_all_gather (parallel_context = parallel_context )
149
148
@@ -156,7 +155,7 @@ def _init_model(self, model_config, model):
156
155
return NanotronLightevalModel (
157
156
checkpoint_path = os .path .dirname (self .pipeline_parameters .nanotron_checkpoint_path ),
158
157
nanotron_config = self .model_config ,
159
- parallel_context = self .accelerator ,
158
+ parallel_context = self .parallel_context ,
160
159
debug_one_layer_model = False ,
161
160
model_class = None ,
162
161
env_config = self .pipeline_parameters .env_config ,
0 commit comments