8000 dalle-mini/tools/train/train.py at queue · hoopper/dalle-mini · GitHub
[go: up one dir, main page]
Skip to content

train.py

Latest commit

 

History

History
1636 lines (1481 loc) · 59.9 KB

train.py

File metadata and controls

1636 lines (1481 loc) · 59.9 KB
Raw
default="dp",
metadata={
"help": "Whether to shard the optimizer across data devices (dp), model devices (mp) or both (2d)."
},
)
num_train_epochs: int = field(
default=3, metadata={"help": "Total number of training epochs to perform."}
)
warmup_steps: int = field(
default=0, metadata={"help": "Linear warmup over warmup_steps."}
)
lr_decay: str = field(
default=None,
metadata={
"help": "Decay to be used in the learning rate scheduler. Can be None (default), linear or exponential."
},
)
lr_transition_steps: int = field(
default=None,
metadata={
"help": "Number of transition steps associated with learning rate decay when using exponential decay."
},
)
lr_decay_rate: float = field(
default=None,
metadata={
"help": "Decay rate associated with learning rate when using exponential decay."
},
)
lr_staircase: bool = field(
default=False,
metadata={
"help": "Whether to use staircase or continuous learning rate when using exponential decay."
},
)
lr_offset: int = field(
default=0,
metadata={"help": "Number of steps to offset learning rate and keep it at 0."},
)
logging_steps: int = field(
default=40, metadata={"help": "Log every X updates steps."}
)
eval_steps: int = field(
default=400, metadata={"help": "Run an evaluation every X steps."}
)
save_steps: int = field(
default=4000, metadata={"help": "Save checkpoint every X updates steps."}
)
log_model: bool = field(
default=False,
metadata={"help": "Log model to wandb at `save_steps` frequency."},
)
log_norm_steps: int = field(
default=True,
metadata={"help": "Log parameters and gradients norm at this frequency."},
)
log_histogram_steps: int = field(
default=False,
metadata={
"help": "Log parameters and gradients histograms at this frequency. Slows down training."
},
)
seed_model: int = field(
default=42,
metadata={
"help": "Random seed for the model that will be set at the beginning of training."
},
)
wandb_entity: Optional[str] = field(
default=None,
metadata={"help": "The wandb entity to use (for teams)."},
)
wandb_project: str = field(
default="dalle-mini",
metadata={"help": "The name of the wandb project."},
)
wandb_job_type: str = field(
default="Seq2Seq",
metadata={"help": "The name of the wandb job type."},
)
assert_TPU_available: bool = field(
default=False,
metadata={"help": "Verify that TPU is not in use."},
)
use_vmap_trick: bool = field(
default=True,
metadata={"help": "Verify that TPU is not in use."},
)
mp_devices: Optional[int] = field(
default=1,
metadata={
"help": "Number of devices required for model parallelism. The other dimension of available devices is used for data parallelism."
},
)
dp_devices: int = field(init=False)
def __post_init__(self):
if self.assert_TPU_available:
assert (
jax.local_device_count() == 8
), "TPUs in use, please check running processes"
if self.output_dir.startswith("gs://"):
assert (
storage is not None
), 'Could not find google.storage. Install with "pip install google-cloud-storage"'
assert self.optim in [
"distributed_shampoo",
"adam",
"adafactor",
], f"Selected optimizer not supported: {self.optim}"
if self.optim == "adafactor" and self.weight_decay == 0:
self.weight_decay = None
assert self.graft_type in [
"rmsprop_normalized",
"rmsprop",
"adagrad",
"adagrad_normalized",
"sgd",
"sqrt_n",
], f"Selected graft type not supported: {self.graft_type}"
assert self.lr_decay in [
None,
"linear",
"exponential",
], f"Selected learning rate decay not supported: {self.lr_decay}"
if self.per_device_eval_batch_size is None:
self.per_device_eval_batch_size = self.per_device_train_batch_size
if self.log_norm_steps is True:
self.log_norm_steps = self.logging_steps
if not self.do_train:
self.num_train_epochs = 1
if (
os.path.exists(self.output_dir)
and os.listdir(self.output_dir)
and self.do_train
and not self.overwrite_output_dir
):
raise ValueError(
f"Output directory ({self.output_dir}) already exists and is not empty."
"Use --overwrite_output_dir to overcome."
)
assert self.shard_shampoo_across in [
"dp",
"mp",
"2d",
], f"Shard shampoo across {self.shard_shampoo_across} not supported."
assert (
self.mp_devices > 0
), f"Number of devices for model parallelism must be > 0"
assert (
jax.device_count() % self.mp_devices == 0
), f"Number of available devices ({jax.device_count()} must be divisible by number of devices used for model parallelism ({self.mp_devices})."
self.dp_devices = jax.device_count() // self.mp_devices
def split_params(data):
"""Split params between scanned and non-scanned"""
flat = traverse_util.flatten_dict(unfreeze(data))
split = {"standard": {}, "scanned_encoder": {}, "scanned_decoder": {}}
for k, v in flat.items():
if "FlaxBartEncoderLayers" in k:
split["scanned_encoder"][k] = v
elif "FlaxBartDecoderLayers" in k:
split["scanned_decoder"][k] = v
else:
split["standard"][k] = v
# remove empty keys
split = {k: v for k, v in split.items() if v}
for k, v in split.items():
split[k] = freeze(traverse_util.unflatten_dict(v))
return split
def unsplit_params(data):
flat = {}
for k in ["standard", "scanned_encoder", "scanned_decoder"]:
if k in data:
flat.update(traverse_util.flatten_dict(unfreeze(data[k])))
return freeze(traverse_util.unflatten_dict(flat))
class TrainState(struct.PyTreeNode):
step: int
params: core.FrozenDict[str, Any]
opt_state: optax.OptState
apply_fn: Callable = struct.field(pytree_node=False)
tx: optax.GradientTransformation = struct.field(pytree_node=False)
dropout_rng: jnp.ndarray = None
epoch: int = 0
train_time: float = 0.0 # total time the model trained
train_samples: int = 0 # number of samples seen
def apply_gradients(self, *, grads, **kwargs):
grads = split_params(grads)
params = split_params(self.params)
opt_state = {}
# we loop over keys: "standard", "scanned_encoder", "scanned_decoder"
for k, param in params.items():
update_fn = self.tx[k].update
if "scanned" in k:
update_fn = jax.vmap(update_fn, in_axes=(0, 0, 0), out_axes=(0, 0))
updates, new_opt_state = update_fn(grads[k], self.opt_state[k], param)
params[k] = optax.apply_updates(param, updates)
opt_state[k] = new_opt_state
params = unsplit_params(params)
return self.replace(
step=self.step + 1,
params=params,
opt_state=freeze(opt_state),
**kwargs,
)
@classmethod
def create(cls, *, apply_fn, params, tx, **kwargs):
opt_state = {}
for k, p in split_params(params).items():
init_fn = tx[k].init
if "scanned" in k:
init_fn = jax.vmap(init_fn)
opt_state[k] = init_fn(p)
return cls(
step=0,
apply_fn=apply_fn,
params=params,
tx=tx,
opt_state=freeze(opt_state),
**kwargs,
)
def main():
# See all possible arguments by passing the --help flag to this script.
parser = HfArgumentParser(
(ModelArguments, DataTrainingArguments, TrainingArguments)
)
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
# If we pass only one argument to the script and it's the path to a json file,
# let's parse it to get our arguments.
model_args, data_args, training_args = parser.parse_json_file(
json_file=os.path.abspath(sys.argv[1])
)
else:
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
# check arguments
if training_args.mp_devices > jax.local_device_count():
assert (
data_args.seed_dataset is not None
), "Seed dataset must be provided when model is split over multiple hosts"
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
level=logging.INFO,
)
# Setup logging, we only want one process per machine to log things on the screen.
logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
if jax.process_index() == 0:
datasets.utils.logging.set_verbosity_warning()
transformers.utils.logging.set_verbosity_info()
else:
datasets.utils.logging.set_verbosity_error()
transformers.utils.logging.set_verbosity_error()
# Set the verbosity to info of the Transformers logger (on main process only):
logger.info(f"Training/evaluation parameters {training_args}")
# Load dataset
dataset = Dataset(
**asdict(data_args),
do_train=training_args.do_train,
do_eval=training_args.do_eval,
)
logger.info(f"Local TPUs: {jax.local_device_count()}")
logger.info(f"Global TPUs: {jax.device_count()}")
# Set up wandb run
if jax.process_index() == 0:
wandb.init(
entity=training_args.wandb_entity,
project=training_args.wandb_project,
job_type=training_args.wandb_job_type,
config=parser.parse_args(),
)
# Set up our new model config
config_args = {
k: getattr(model_args, k)
for k in ["dropout", "activation_dropout", "attention_dropout"]
if getattr(model_args, k) is not None
}
if model_args.config_name:
config = DalleBartConfig.from_pretrained(model_args.config_name)
config.gradient_checkpointing = training_args.gradient_checkpointing
for k, v in config_args.items():
setattr(config, k, v)
else:
config = None
# Load or create new model
if model_args.model_name_or_path:
model, params = DalleBart.from_pretrained(
model_args.model_name_or_path,
config=config,
seed=training_args.seed_model,
dtype=getattr(jnp, model_args.dtype),
_do_init=False, # we overwrite them with loaded checkpoint
gradient_checkpointing=training_args.gradient_checkpointing,
**config_args,
)
else:
model = DalleBart(
config,
seed=training_args.seed_model,
dtype=getattr(jnp, model_args.dtype),
_do_init=False,
)
params = None
params_shape = model.params_shape_tree
# get model metadata
model_metadata = model_args.get_metadata()
# get PartitionSpec for model params (required to be a dict)
param_spec = set_partitions(params_shape, model.config.use_scan)
params_shape = freeze(params_shape)
if params is not None:
params = freeze(params)
# Load tokenizer
tokenizer = DalleBartTokenizer.from_pretrained(
model_args.tokenizer_name, use_fast=True
)
# Preprocessing the datasets.
# We need to normalize and tokenize inputs and targets.
dataset.preprocess(tokenizer=tokenizer, config=model.config)
# Initialize our training
dropout_rng = jax.random.PRNGKey(training_args.seed_model)
# Store some constant
num_epochs = training_args.num_train_epochs
# batch size
batch_size_per_node_per_grad_step = (
training_args.per_device_train_batch_size
* jax.local_device_count()
// training_args.mp_devices
)
batch_size_per_node = (
batch_size_per_node_per_grad_step * training_args.gradient_accumulation_steps
)
batch_size_per_step = batch_size_per_node * jax.process_count()
eval_batch_size_per_node = (
training_args.per_device_eval_batch_size
* jax.local_device_count()
// training_args.mp_devices
)
eval_batch_size_per_step = eval_batch_size_per_node * jax.process_count()
len_train_dataset, len_eval_dataset = dataset.length
steps_per_epoch = (
len_train_dataset // batch_size_per_node
if len_train_dataset is not None
else None
)
num_train_steps = (
steps_per_epoch * num_epochs if steps_per_epoch is not None else None
)
num_params = model.num_params(params_shape)
logger.info("***** Running training *****")
logger.info(f" Num examples = {len_train_dataset}")
logger.info(f" Num Epochs = {num_epochs}")
logger.info(
f" Batch size per dp device = {training_args.per_device_train_batch_size}"
)
logger.info(f" Number of devices = {jax.device_count()}")
logger.info(
f" Gradient accumulation steps = {training_args.gradient_accumulation_steps}"
)
logger.info(f" Batch size per update = {batch_size_per_step}")
logger.info(f" Model parameters = {num_params:,}")
# set up wandb run
if jax.process_index() == 0:
# set default x-axis as 'train/step'
wandb.define_metric("*", step_metric="train/step")
# add interesting config parameters
wandb.config.update(
{
"len_train_dataset": len_train_dataset,
"len_eval_dataset": len_eval_dataset,
"batch_size_per_step": batch_size_per_step,
"num_params": num_params,
"model_config": model.config.to_dict(),
"num_devices": jax.device_count(),
"versions": {
"jax": jax.__version__,
"jaxlib": jaxlib.__version__,
"flax": flax.__version__,
"transformers": transformers.__version__,
"datasets": datasets.__version__,
"wandb": wandb.__version__,
"dalle_mini": dalle_mini.__version__,
},
}
)
# Create learning rate schedule
def create_learning_rate_fn() -> Callable[[int], jnp.array]:
"""Create the learning rate function."""
warmup_fn = optax.linear_schedule(
init_value=0.0,
end_value=training_args.learning_rate,
transition_steps=training_args.warmup_steps + 1, # ensure not 0
)
last_boundary = training_args.warmup_steps
# offset step when resuming
if training_args.lr_offset:
warmup_fn = optax.join_schedules(
schedules=[optax.constant_schedule(0.0), warmup_fn],
boundaries=[training_args.lr_offset],
)
last_boundary += training_args.lr_offset
if training_args.lr_decay is None:
return warmup_fn
elif training_args.lr_decay == "linear":
assert (
num_train_steps is not None
), "linear decay requires knowing the dataset length"
decay_fn = optax.linear_schedule(
init_value=training_args.learning_rate,
end_value=0,
transition_steps=num_train_steps - training_args.warmup_steps,
)
elif training_args.lr_decay == "exponential":
decay_fn = optax.exponential_decay(
init_value=training_args.learning_rate,
transition_steps=training_args.lr_transition_steps,
decay_rate=training_args.lr_decay_rate,
staircase=training_args.lr_staircase,
)
schedule_fn = optax.join_schedules(
schedules=[warmup_fn, decay_fn],
boundaries=[last_boundary],
)
return schedule_fn
learning_rate_fn = create_learning_rate_fn()
# create adam optimizer
if training_args.optim == "distributed_shampoo":
# parameters from https://github.com/tensorflow/lingvo/blob/03ee9d7cd50764b0424c7c863733c91fc0b053ec/lingvo/jax/optimizers.py#L729
graft_type = {
"sgd": GraftingType.SGD,
"adagrad": GraftingType.ADAGRAD,
"rmsprop": GraftingType.RMSPROP,
"rmsprop_normalized": GraftingType.RMSPROP_NORMALIZED,
"sqrt_n": GraftingType.SQRT_N,
"adagrad_normalized": GraftingType.ADAGRAD_NORMALIZED,
}[training_args.graft_type]
statistics_partition_spec = (
PartitionSpec(None, training_args.shard_shampoo_across, None)
if training_args.shard_shampoo_across != "2d"
else PartitionSpec(None, "dp", "mp")
)
opt = distributed_shampoo(
learning_rate_fn,
block_size=training_args.block_size,
beta1=training_args.beta1,
beta2=training_args.beta2,
diagonal_epsilon=1e-10,
matrix_epsilon=1e-6,
weight_decay=training_args.weight_decay,
start_preconditioning_step=max(
training_args.preconditioning_compute_steps + 1, 101
),
preconditioning_compute_steps=training_args.preconditioning_compute_steps,
statistics_compute_steps=1,
best_effort_shape_interpretation=True,
graft_type=graft_type,
nesterov=training_args.nesterov,
exponent_override=0,
statistics_partition_spec=statistics_partition_spec,
preconditioner_partition_spec=PartitionSpec(
training_args.shard_shampoo_across, None, None
)
if training_args.shard_shampoo_across != "2d"
else PartitionSpec(
"mp" if training_args.mp_devices > training_args.dp_devices else "dp",
None,
None,
),
num_devices_for_pjit=training_args.dp_devices,
shard_optimizer_states=True,
inverse_failure_threshold=0.1,
moving_average_for_momentum=True,
skip_preconditioning_dim_size_gt=training_args.skip_preconditioning_dim_size_gt,
clip_by_scaled_gradient_norm=None,
precision=jax.lax.Precision.HIGHEST,
best_effort_memory_usage_reduction=training_args.optim_quantized,
)
# get the real optimizer and helper functions
update_fn = opt.update
optimizer = {}
opt_fn = {}
for k, p in split_params(params_shape).items():
if "scanned" in k:
p = jax.eval_shape(lambda x: jax.tree_map(lambda y: y[0], x), p)
optimizer[k] = opt.init(p)
opt_fn[k] = NamedTuple("opt_fn", pspec_fn=Any, shape_and_dtype_fn=Any)(
optimizer[k].pspec_fn, optimizer[k].shape_and_dtype_fn
)
optimizer[k] = optax.GradientTransformation(optimizer[k].init_fn, update_fn)
elif training_args.optim == "adam":
optimizer = optax.adamw(
learning_rate=learning_rate_fn,
b1=training_args.beta1,
b2=training_args.beta2,
eps=training_args.adam_epsilon,
weight_decay=training_args.weight_decay,
)
optimizer = {k: optimizer for k in split_params(params_shape)}
elif training_args.optim == "adafactor":
# We use the default parameters here to initialize adafactor,
# For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
optimizer = optax.adafactor(
learning_rate=learning_rate_fn,
clipping_threshold=training_args.max_grad_norm,
weight_decay_rate=training_args.weight_decay,
)
optimizer = {k: optimizer for k in split_params(params_shape)}
# get PartitionSpec for optimizer state
def get_opt_state_spec_and_shape():
# get opt_state shape without actual init
opt_state_shape = {}
for k, p in split_params(params_shape).items():
if "scanned" not in k:
opt_state_shape[k] = jax.eval_shape(optimizer[k].init, p)
else:
opt_state_shape[k] = jax.eval_shape(jax.vmap(optimizer[k].init), p)
if training_args.optim == "adafactor":
# factorized state must be replicated (rank different than params)
opt_state_spec = {k: None for k in split_params(params_shape)}
elif training_args.optim in ["adam", "distributed_shampoo"]:
def _opt_state_spec_per_leaf(x, spec):
if isinstance(x, FrozenDict):
# variables with same structure as params
return spec
else:
# other variables such as count
return None
split_spec = split_params(set_partitions(params_shape, False))
opt_state_spec = {}
for k, p in split_params(params_shape).items():
if "scanned" in k:
p = jax.eval_shape(lambda x: jax.tree_map(lambda y: y[0], x), p)
if training_args.optim == "adam":
opt_state_spec[k] = jax.tree_map(
_opt_state_spec_per_leaf,
opt_state_shape[k],
split_spec[k],
# return None spec for empty elements
is_leaf=lambda x: isinstance(x, (FrozenDict, optax.EmptyState)),
)
elif training_args.optim == "distributed_shampoo":
opt_state_spec[k] = opt_fn[k].pspec_fn(
p,
split_spec[k],
statistics_partition_spec,
)
# add dimension for scanned params
if "scanned" in k:
opt_state_spec[k] = jax.tree_map(
lambda x: PartitionSpec(*(None,) + x)
if x is not None
else None,
opt_state_spec[k],
is_leaf=lambda x: isinstance(x, PartitionSpec),
)
else:
raise NotImplementedError
return freeze(opt_state_spec), freeze(opt_state_shape)
opt_state_spec, opt_state_shape = get_opt_state_spec_and_shape()
# create a mesh
mesh_shape = (training_args.dp_devices, training_args.mp_devices)
devices = np.asarray(jax.devices()).reshape(*mesh_shape)
View remainder of file in raw view
0