pixtral SFT

[shruthan]

✨ Description

Some bug fixes for Image+Text SFTs

🔍 Type of change

Select all that apply:

• 🐛 Bug fix (non-breaking change that addresses a specific issue)
• 🚀 New feature (non-breaking change that adds functionality)
• ⚠️ Breaking change (a change that could affect existing functionality)
• 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• 📝 Documentation change (updates documentation, including new content or typo fixes)
• 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

📝 Changes

List the key changes introduced in this PR:

1. Change A
2. Change B

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General

• 📜 I have read and followed the contributing guidelines.
• 🏷️ I am using a clear and descriptive PR title that summarizes the key change or feature introduced.
• 🎉 The functionality is complete, and I have tested the changes.
• 📝 I have updated the documentation if needed.
• ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
• 🧩 I have commented my code, especially in hard-to-understand areas.

Dependencies and Configuration

• 🐋 I have updated the Docker configuration or dependencies, if applicable.
• 🔄 I have ensured compatibility with the existing setup after dependency changes.

Testing

• 🧪 I have added or updated tests to cover my changes.
• ✔️ New and existing tests pass locally with my changes.
• 🚦 I have tested these changes on GPUs and verified training stability.
• 🏋️ I have tested the changes on realistic training workloads, if applicable.

Performance Impact

• 📊 I have run benchmarks where applicable to evaluate the performance impact.
• ✅ The benchmarks show no performance regression.
• 🚀 The benchmarks indicate a potential performance improvement.
• ⚠️ The benchmarks indicate a potential performance degradation.
• 📈 I have provided benchmark results and detailed any performance impact below, if applicable.

📊 Performance Impact Details

If there is any impact on performance, describe it and provide benchmark results, if applicable:

---

🗒️ Additional Notes

Include any additional context, information, or considerations here, such as known issues, follow-up tasks, or backward compatibility concerns.

[sohamparikh]

it's already read above, why read it here?

[shruthan]

There's another flag written after images, this is to read that, but not sure what the assignment should be

Fast-LLM/fast_llm/data/dataset/gpt/memmap.py

Lines 407 to 412 in 2f85615

	# Placeholder flag for preference spans
	idx_stream.write(struct.pack("<B", 0))
	# Flag to indicate whether images are present
	idx_stream.write(struct.pack("<B", 1 if total_images > 0 else 0))
	# Flag to indicate whether preference loss-masking spans are present
	idx_stream.write(struct.pack("<B", 1 if chosen_spans.size > 0 and rejected_spans.size > 0 else 0))

[sohamparikh]

oh yeah that order should be flipped, the chosen_spans byte should be before total_images, i'll fix it in my branch

[RaymondLi0]

This would break files with version==3 right?

It seems to me that we should rather fix the order in which those flags are dumped in the idx file below

[sohamparikh]

@RaymondLi0 yes, I'm planning to fix it in #227

[sohamparikh]

Suggested change

	total_pixels_needed = sum(
	length[0] * length[1] * 3 for length in self._image_lengths[idx]
	)
	total_pixels_needed = self._image_lengths[idx].prod(initial=3, axis=1).sum()

[sohamparikh]

can leave it as using .prod?

[sohamparikh]

use a bool has_images = bool(sample.image_positions) and use it below as well?

[shruthan]

I think bool(sample.image_positions) will throw ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all() when there are more than one image_positions: bool(np.array([2, 3]))

[sohamparikh]

or simply has_images = True if sample.image_positions else False?

[shruthan]

if sample.image_positions would throw the same error, changed it to has_image_positions = sample.image_positions is not None

[sohamparikh]

LGTM! Thanks!